### A Pluto.jl notebook ###
# v0.14.6

using Markdown
using InteractiveUtils

# ╔═╡ 755e081d-7176-4325-b481-d5d67db6176c
begin
    import Pkg
    Pkg.activate(mktempdir())
    Pkg.add([
        Pkg.PackageSpec(name="Plots", version="1"),
        Pkg.PackageSpec(name="PlutoUI", version="0.7"),
	    Pkg.PackageSpec(name="DifferentialEquations", version="6"),
	    Pkg.PackageSpec(name="Parameters", version="0.12"),
        Pkg.PackageSpec(name="LabelledArrays", version="1"),
        Pkg.PackageSpec(name="Setfield", version="0.7"),
    ])
    using Plots, PlutoUI, LinearAlgebra, DifferentialEquations, Parameters, LabelledArrays, Setfield

	Plots.gr(fmt=:png, lw=2)

	PlutoUI.TableOfContents()
end

# ╔═╡ 4a209b69-b1b4-4a63-8786-c2b2a1a9b6ee
md"""
# Figure 4.1 to 4.5

## Figure 4.1, 4.2, and 4.3

Steady states and phase plots in an assymetric network
"""

# ╔═╡ 9b0868de-687e-44f0-a3b8-2321cebd3a52
params = (K1=20.0, K2=5.0, K3=5.0, K4=5.0, K5=2.0, N=4)

# ╔═╡ 69f0e704-2a02-430d-97fd-4a7e5b1fc031
u0s = (LVector(A=0.0, B=0.0),
       LVector(A=0.5, B=0.6),
       LVector(A=0.17, B=1.1),
       LVector(A=0.25, B=1.9),
       LVector(A=1.85, B=1.70))

# ╔═╡ 298447c2-7e64-422c-9d25-9fb4dee35441
md"""
## Figure 4.4, 4.5

Vector fields
"""

# ╔═╡ 0b3412c0-bde3-11eb-0869-17548acad535
md"""
## Appendix
Running environment and some auxillary functions.

**You can use [this helper tool](https://fonsp.com/article-test-3/pkghelper.html) to generate these commands!**
"""

# ╔═╡ 3d3f9ef7-4195-413a-a4d4-563f327ccda4
begin
	# Convenience functions
    hill(x, k) = x / (x + k)
    hill(x, k, n) = hill(x^n, k^n)
    exprel(x) = ifelse(x≈zero(x), one(x), x / expm1(x))
end

# ╔═╡ f8ff273a-e1c7-4f00-9f9d-0ecb99f90b1e
function model!(du, u, p, t)
    @unpack K1, K2, K3, K4, K5, N = p
    @unpack A, B = u

    v1 = K1 * hill(1, B, N)
    v5 = K5 * A

    du.A = v1 - v5 - K3 * A
    du.B = K2 + v5 - K4 * B
end

# ╔═╡ 632c9de0-dc67-4a7f-882d-e04403d24d0e
sols = map(u0 -> solve(ODEProblem(model!, u0, 1.5, params)), u0s);

# ╔═╡ 8d6fb7a3-b5b0-4d49-9b07-8dcb1e2105cc
plot(sols[1], xlabel="Time", ylabel="Concentration", title="Fig. 4.2A (Time series)")

# ╔═╡ afed1e2a-40a6-4ddf-8988-6a66661e5063
plot(sols[1], vars=(1, 2), xlabel="[A]", ylabel="[B]", aspect_ratio=:equal,
     title="Fig. 4.2B (Phase plot)", ylims=(0.0, 2.0), xlims=(0.0, 2.0),
     legend=nothing)

# ╔═╡ ff7eed37-4569-45c1-af5c-aebda8fb2003
begin
	p3 = plot()

	for sol in sols
		plot!(p3, sol, linealpha=0.5, legend = nothing)
	end

	plot!(p3, xlabel="Time", ylabel="Concentration", title="Fig. 4.3A (Time series)")
end

# ╔═╡ d9295b32-da22-45e6-910b-6bd1126bb92b
begin
	p4 = plot()

	for sol in sols
		plot!(p4, sol, vars=(1, 2), linealpha=0.7, legend = nothing)
	end

	plot!(p4, aspect_ratio=:equal, title="Fig. 4.3B (Phase plot)", xlabel="[A]", ylabel="[B]", ylims=(0.0, 2.0), xlims=(0.0, 2.0), size=(600, 600))
end

# ╔═╡ a243b6e3-dd6b-47d0-b807-db88622bb3a4
# vector field
function df(x, y, p = params, scale=20)
	u = LVector(A=x, B=y)
	du = similar(u)
	model!(du, u, p, 0.0)
	return du ./ (norm(du)^0.5 * scale)
end

# ╔═╡ 0dd10ca8-71eb-4f68-89d7-a4717e06afc1
begin
	xx = [x for y in 0.0:0.1:2.0, x in 0.0:0.1:2.0]
	yy = [y for y in 0.0:0.1:2.0, x in 0.0:0.1:2.0]
	p1 = quiver(xx, yy, quiver=df, line=(:lightblue))

	for sol in sols
		plot!(p1, sol, vars=(1, 2), legend = nothing, line=(:green))
	end

	plot!(p1, aspect_ratio=:equal, title="Fig. 4.4A (Phase Plot with vector field)",
      xlabel="[A]", ylabel="[B]", xlim=(0.0, 2.0), ylim=(0.0, 2.0), size=(600, 600))
	p1
end

# ╔═╡ 799cb2ff-a7a2-4947-99c2-d23c06f0b920
# Nullclines
begin
	nullcline_a(b, p) = p.K1 / (p.K5 + p.K4)  * hill(1, b, p.N)
	nullcline_a(b) = nullcline_a(b, params)
	nullcline_b(b, p) = (p.K4*b - p.K2) / p.K5
	nullcline_b(b) = nullcline_b(b, params)
end

# ╔═╡ 160b9f34-e089-40b1-8976-b1a168da2044
begin
	# Figure 4.5A
	p45a = plot(aspect_ratio=:equal, title="Fig. 4.5A, Phase plot with nullclines")

	# Phase plots
	for sol in sols
		plot!(p45a, sol, vars=(1, 2), linealpha=0.7, lab=nothing)
	end

	# Parametric plotting for nullcline
	plot!(p45a, nullcline_a, identity, 0.0, 2.0, label="A nullcline", line=(:black, :dot))
	plot!(p45a, nullcline_b, identity, 0.0, 2.0, label="B nullcline", line=(:black, :dash))
	plot!(p45a, xlim=(0.0, 2.0), ylim=(0.0, 2.0), legend=:bottomright, size=(600, 600), xlabel="[A]", ylabel="[B]")

	p45a
end

# ╔═╡ 30ecb42b-ed42-455a-a933-2f00a4495907
begin
	p45b = quiver(xx, yy, quiver=df, line=(:lightblue), title="Fig. 4.5B, Vector field with nullclines", xlabel="[A]", ylabel="[B]")
	plot!(p45b, nullcline_a, identity, 0.0, 2.0, label="A nullcline", line=(:black, :dot))
	plot!(p45b, nullcline_b, identity, 0.0, 2.0, label="B nullcline", line=(:black, :dash))
	plot!(p45b, aspect_ratio=1.0, xlim=(0.0, 2.0), ylim=(0.0, 2.0), legend=:bottomright, size=(600, 600))

	p45b
end

# ╔═╡ Cell order:
# ╠═4a209b69-b1b4-4a63-8786-c2b2a1a9b6ee
# ╠═f8ff273a-e1c7-4f00-9f9d-0ecb99f90b1e
# ╠═9b0868de-687e-44f0-a3b8-2321cebd3a52
# ╠═69f0e704-2a02-430d-97fd-4a7e5b1fc031
# ╠═632c9de0-dc67-4a7f-882d-e04403d24d0e
# ╠═8d6fb7a3-b5b0-4d49-9b07-8dcb1e2105cc
# ╠═afed1e2a-40a6-4ddf-8988-6a66661e5063
# ╠═ff7eed37-4569-45c1-af5c-aebda8fb2003
# ╠═d9295b32-da22-45e6-910b-6bd1126bb92b
# ╠═298447c2-7e64-422c-9d25-9fb4dee35441
# ╠═799cb2ff-a7a2-4947-99c2-d23c06f0b920
# ╠═a243b6e3-dd6b-47d0-b807-db88622bb3a4
# ╠═0dd10ca8-71eb-4f68-89d7-a4717e06afc1
# ╠═160b9f34-e089-40b1-8976-b1a168da2044
# ╠═30ecb42b-ed42-455a-a933-2f00a4495907
# ╠═0b3412c0-bde3-11eb-0869-17548acad535
# ╠═3d3f9ef7-4195-413a-a4d4-563f327ccda4
# ╠═755e081d-7176-4325-b481-d5d67db6176c


"""
Tools for the exploration of chaos and nonlinear dynamics
"""
module ChaosTools

using Reexport
@reexport using DynamicalSystemsBase

using DynamicalSystemsBase: DS, DDS, CDS
using DynamicalSystemsBase: MDI, TDI
using DynamicalSystemsBase: stateeltype
using DynamicalSystemsBase.DiffEqBase: AbstractODEIntegrator, u_modified!, DEIntegrator
DEI = DEIntegrator

include("orbitdiagrams/discrete_diagram.jl")
include("orbitdiagrams/poincare.jl")

include("dimensions/entropies.jl")
include("dimensions/dims.jl")
include("dimensions/correlationdim.jl")

include("nlts.jl")

include("period_return/periodic_points.jl")
include("period_return/period.jl")
include("period_return/transit_times_statistics.jl")

include("chaosdetection/lyapunovs.jl")
include("chaosdetection/gali.jl")
include("chaosdetection/expansionentropy.jl")
include("chaosdetection/partially_predictable.jl")
include("chaosdetection/testchaos01.jl")

# Ugly methods that shouldn't exist:
include("ugliness.jl")

display_update = false
version = "1.10.0"
update_name = "update_v$version"

if display_update
if !isfile(joinpath(@__DIR__, update_name))
printstyled(stdout,
"""
\nUpdate message: ChaosTools v$version

A method to calculate the expansion entropy for discrete and
continuous systems is now included as `expansionentropy`!

See B. Hunt & E. Ott, ‘Defining Chaos’, Chaos 25.9 (2015).
\n
"""; color = :light_magenta)
touch(joinpath(@__DIR__, update_name))
end
end

end # module


#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/build_literate/Analyze_NMR_IR_Data.ipynb)

# # Inversion Recovery Experiments
# The following code replicates the NMR data analysis in Figs. 4-6 and complements the paper with additional analyses that are not shown in the paper in the interest of brevity.

# For this analysis we need the following packages:
using MRIgeneralizedBloch
using DifferentialEquations
using LinearAlgebra
using LsqFit
using Statistics
import Pingouin
using Printf
using Formatting
using Plots
plotlyjs(bg = RGBA(31/255,36/255,36/255,1.0), ticks=:native); #hide #!nb

# The raw data is stored in a separate [github repository](https://github.com/JakobAsslaender/MRIgeneralizedBloch_NMRData) and the following functions return the URL to the individual files:
MnCl2_data(TRF_scale) = string("https://github.com/JakobAsslaender/MRIgeneralizedBloch_NMRData/blob/main/20210419_1mM_MnCl2/ja_IR_v2%20(", TRF_scale, ")/1/data.2d?raw=true")
BSA_data(TRF_scale)   = string("https://github.com/JakobAsslaender/MRIgeneralizedBloch_NMRData/blob/main/20210416_15%25BSA_2ndBatch/ja_IR_v2%20(", TRF_scale, ")/1/data.2d?raw=true");

# which can be loaded with functions implemented in this file:
include(string(pathof(MRIgeneralizedBloch), "/../../docs/src/load_NMR_data.jl"));

# ## MnCl``_2`` Sample
# ### ``T_2^{*,f}`` Estimation
# We estimate ``T_2^{*,f}`` by fitting a mono-exponential decay curve to the FID of the acquisition with ``T_\text{RF} = 22.8``μs and ``T_\text{i} = 5``s.
M = load_Data(MnCl2_data(1))
M = M[:,1]; # select Tᵢ = 5s

# The data was measured at the following time points in units of seconds:
T_dwell = 100e-6 # s
TE = T_dwell * ((1:length(M)) .+ 7) # s

# Note that the signal is an FID, so the phrase *echo* time is a bit misleading.

# The function [curve_fit](https://julianlsolvers.github.io/LsqFit.jl/latest/api/#LsqFit.curve_fit) from the [LsqFit.jl](https://julianlsolvers.github.io/LsqFit.jl/latest/) package is only implemented for real-valued models. To accommodate this, we need to split the data into its real and imaginary part:
TEreal = [TE;TE]
Mreal = [real(M);imag(M)];

# Here, we are using a simple mono-exponential model with a complex-valued scaling factor `p[1] + 1im p[2]`, the decay time ``T_2^{*,f} =`` `p[3]`, and the Larmor frequency `p[4]`:
FID_model(t, p) = @. [p[1] * exp(- t[1:end ÷ 2] / p[3]) * cos(p[4] * t[1:end ÷ 2]); p[2] * exp(- t[end ÷ 2 + 1:end] / p[3]) * sin(p[4] * t[end ÷ 2 + 1:end])];

# Fitting this model to the NMR data estimates ``T_2^{*,f}``:
fit = curve_fit(FID_model, TEreal, Mreal, [1, 1, 0.1, 0])
T₂star_MnCl2 = fit.param[3] # s
# seconds and its uncertainty (also in units of seconds)
stderror(fit)[3] # s

# Visually, the plot and the data show good agreement:
Mfitted = FID_model(TEreal, fit.param)
Mfitted = Mfitted[1:end÷2] + 1im * Mfitted[end÷2+1:end]
p = plot(xlabel="TE [s]", ylabel="|FID(TE)| [a.u.]")
plot!(p, TE, abs.(M), label="data")
plot!(p, TE, abs.(Mfitted), label=@sprintf("fit with T₂* = %2.3f ms", 1e3 * T₂star_MnCl2))
#md Main.HTMLPlot(p) #hide

# The relative residual norm of the fit, i.e. ``||\text{residual}||_2/||M||_2`` is
norm(fit.resid) / norm(M)

# Despite its small ``\ell_2``-norm, the Shapiro-Wilk test indicates that the residual is not Gaussian or normal distributed at a significance level of `α=0.05`
Pingouin.normality(fit.resid, α=0.05)
# We note that mono-exponential ``T_2^*`` decays assume a Lorentzian distributed magnetic field, which is in general an assumption rather than a well-founded theory.

# ### Mono-Exponential IR Model
# We performed several experiments in which we inverted the thermal equilibrium magnetization with rectangular π-pulses with the following pulse durations (in seconds):
Tʳᶠmin = 22.8e-6 # s - shortest Tʳᶠ possible on the NMR
TRF_scale = [1;2;5:5:40] # scaling factor
Tʳᶠ = TRF_scale * Tʳᶠmin # s

# and acquired inversion recovery data at exponentially spaced inversion times (in seconds):
Tᵢ = exp.(range(log(3e-3), log(5), length=20)) # s
Tᵢ .+= 12 * Tʳᶠmin + (13 * 15.065 - 5) * 1e-6 # s - correction factors

# We calculate the Rabi frequencies of the RF pulses and a finer grid of ``T_\text{i}`` to plot the IR model:
ω₁ = π ./ Tʳᶠ # rad/s
Tᵢplot = exp.(range(log(Tᵢ[1]), log(Tᵢ[end]), length=500)); # s

# After loading and normalizing the data
M = zeros(Float64, length(Tᵢ), length(TRF_scale))
for i = 1:length(TRF_scale)
    M[:,i] = load_first_datapoint(MnCl2_data(TRF_scale[i]))
end
M ./= maximum(M);

#src #############################################################################
#src # export MnCl2 data
#src #############################################################################
io = open(expanduser(string("~/Documents/Paper/2021_MT_IDE/Figures/IR_data_MnCl2.txt")), "w") #src
write(io, "TI_s") #src
for i = 1:length(Tʳᶠ) #src
    write(io, " z_$(@sprintf("%.2e", Tʳᶠ[i]))") #src
end #src
write(io, " \n")  #src

for j = 1:length(Tᵢ) #src
    write(io, "$(@sprintf("%.2e", Tᵢ[j])) ") #src
    for i = 1:length(Tʳᶠ) #src
        write(io, "$(@sprintf("%.2e", M[j,i])) ") #src
    end #src
    write(io, " \n") #src
end #src
close(io) #src
#src #############################################################################

# we analyze each inversion recovery curve that corresponds to a different ``T_\text{RF}`` separately. This allows us to fit a simple mono-exponential model
standard_IR_model(t, p) = @. p[1] - p[3] * exp(- t * p[2]);
# where `p[1]` is the thermal equilibrium magnetization, `p[2]` ``= T_1``, and `p[1] - p[3]` is the magnetization right after the inversion pulse or, equivalently, `Minv = p[1] / p[3] - 1` is the inversion efficiency, which is 1 for an ideal π-pulse and smaller otherwise. The parameters are initialized with
p0 = [1.0, 1.0, 2.0];

# and we can loop over ``T_\text{RF}`` to perform the fits:
param = similar(M[1,:], Vector{Float64}) #src
R₁ = similar(M[1,:])
Minv = similar(R₁)
residual = similar(R₁)
p = plot(xlabel="Tᵢ [s]", ylabel="zᶠ(Tʳᶠ, Tᵢ) [a.u.]")
for i = 1:length(TRF_scale)
    Mi = @view M[:,i]

    fit = curve_fit(standard_IR_model, Tᵢ, Mi, p0)
    param[i] = fit.param #src

    R₁[i] = fit.param[2]
    Minv[i] = fit.param[3] / fit.param[1] - 1

    residual[i] = norm(fit.resid) / norm(Mi)

    scatter!(p, Tᵢ, Mi, label=@sprintf("Tʳᶠ = %1.2es - data", Tʳᶠ[i]), color=i)
    plot!(p, Tᵢplot, standard_IR_model(Tᵢplot, fit.param), label=@sprintf("fit with R₁ = %.3f/s; MInv = %.3f", R₁[i], Minv[i]), color=i)
end
display(p) #!md
#md Main.HTMLPlot(p) #hide

#src #############################################################################
#src # export fitted curves
#src #############################################################################
Mp = [standard_IR_model(Tᵢplot, param[i]) for i=1:length(Tʳᶠ)] #src
io = open(expanduser(string("~/Documents/Paper/2021_MT_IDE/Figures/IR_monoExp_fit_MnCl2.txt")), "w") #src
write(io, "TI_s") #src
for i = 1:length(Tʳᶠ) #src
    write(io, " z_$(@sprintf("%.2e", Tʳᶠ[i]))") #src
end #src
write(io, " \n")  #src

for j = 1:length(Tᵢplot) #src
    write(io, "$(@sprintf("%.2e", Tᵢplot[j])) ") #src
    for i = 1:length(Tʳᶠ) #src
        write(io, "$(@sprintf("%.2e", Mp[i][j])) ") #src
    end #src
    write(io, " \n") #src
end #src
close(io) #src
#src #############################################################################

# Here, the data measured with different ``T_\text{RF}`` are indicated by markers in different colors, and the corresponding fits are the line plots in the same color. The fitted parameters are denoted in the legend. In the paper, we highlight the estimated inversion efficiency and the relaxation rate of the dataset acquired with ``T_\text{RF}=22.8``μs:
Minv[1]
#-
R₁[1] # 1/s
# and of the dataset acquired with ``T_\text{RF}=912``μs:
Minv[end]
#-
R₁[end] # 1/s
# The mean value of all R₁ estimates is
mean(R₁) # 1/s

# 1/s and its standard deviation in units of 1/s is
std(R₁) # 1/s

# The relative residual norm of the fits is on average
mean(residual)

# ### Global IR Fit
# As an alternative to individual fits to the inversion recovery curves with different ``T_\text{RF}``, we can also perform a global fit that accounts for the ``T_2^{*,f}`` decay during the inversion pulse. The model first simulates the ``T_2^{*,f}`` decay during the inversion pulse, followed by ``T_1`` recovery:
function Bloch_IR_model(p, Tʳᶠ, Tᵢ, T2)
    (m0, m0_inv, R₁) = p
    R2 = 1 / T2

    M = zeros(Float64, length(Tᵢ), length(Tʳᶠ))
    for i = 1:length(Tʳᶠ)
        ## simulate inversion pulse
        ω₁ = π / Tʳᶠ[i]
        H = [-R2 -ω₁ 0 ;
              ω₁ -R₁ R₁;
               0   0 0 ]

        m_inv = m0_inv * (exp(H * Tʳᶠ[i]) * [0,1,1])[2]

        ## simulate T1 recovery
        H = [-R₁ R₁*m0;
               0     0]

        for j = 1:length(Tᵢ)
            M[j,i] = m0 * (exp(H .* (Tᵢ[j] - Tʳᶠ[i] / 2)) * [m_inv,1])[1]
        end
    end
    return vec(M)
end;

# We use the previously estimated ``T_2^{*,f}`` value for the fit:
fit = curve_fit((x, p) -> Bloch_IR_model(p, Tʳᶠ, Tᵢ, T₂star_MnCl2), 1:length(M), vec(M), [ 1, .8, 1])

p = plot(xlabel="Tᵢ [s]", ylabel="zᶠ(Tʳᶠ, Tᵢ) [a.u.]")
for i=1:length(Tʳᶠ)
    scatter!(p, Tᵢ, M[:,i], label=@sprintf("Tʳᶠ = %1.2es", Tʳᶠ[i]), color=i)
    plot!(p, Tᵢplot, Bloch_IR_model(fit.param, Tʳᶠ[i], Tᵢplot, T₂star_MnCl2), label=@sprintf("Tʳᶠ = %1.2es", Tʳᶠ[i]), color=i)
end
display(p) #!md
#md Main.HTMLPlot(p) #hide

# With this global fit, we get a very similar relaxation rate in units of 1/s
R₁_MnCl2 = fit.param[3] # 1/s

# with an uncertainty (also in units of 1/s) of
stderror(fit)[3] # 1/s

# Note that the relative residual norm is somewhat increased compared to individual fits to each inversion recovery curve:
norm(fit.resid) / norm(M)

# ## Bovine Serum Albumin Sample
# ### ``T_2^{*,f}`` Estimation
# We repeat the ``T_2^{*,f}`` estimation for the bovine serum albumin (BSA) sample by fitting a mono-exponential decay curve to the FID of the acquisition with ``T_\text{RF} = 22.8``μs and ``T_\text{i} = 5``s.

M = load_Data(BSA_data(1));
M = M[:,1] # select Tᵢ = 5s
Mreal = [real(M);imag(M)]

fit = curve_fit(FID_model, TEreal, Mreal, [1.0, 1.0, .1, 0.0]);

# The estimated ``T_2^{*,f}`` of the BSA sample is
T₂star_BSA = fit.param[3] # s
# seconds with an uncertainty of
stderror(fit)[3] # s
# seconds.

# Visually, the plot and the data align well for the BSA sample, too:
Mfitted = FID_model(TEreal, fit.param)
Mfitted = Mfitted[1:end÷2] + 1im * Mfitted[end÷2+1:end]
p = plot(xlabel="TE [s]", ylabel="|FID(TE)| [a.u.]")
plot!(p, TE, abs.(M), label="data")
plot!(p, TE, abs.(Mfitted), label=@sprintf("fit with T₂* = %2.3f ms", 1e3 * T₂star_BSA))
#md Main.HTMLPlot(p) #hide

# The relative residual norm (``||\text{residual}||_2/||M||_2``) is
norm(fit.resid) / norm(M)

# Despite the small residual, the Shapiro-Wilk test indicates that the residual is not normal distributed for this sample either:
Pingouin.normality(fit.resid, α=0.05)

# ### Mono-Exponential IR Model
# We also fit a mono-exponential model to each inversion recovery curve of the BSA data:
M = zeros(Float64, length(Tᵢ), length(TRF_scale))
for i = 1:length(TRF_scale)
    M[:,i] = load_first_datapoint(BSA_data(TRF_scale[i]))
end
M ./= maximum(M)
#src #############################################################################
#src # export BSA data
#src #############################################################################
io = open(expanduser(string("~/Documents/Paper/2021_MT_IDE/Figures/IR_data_BSA.txt")), "w") #src
write(io, "TI_s") #src
for i = 1:length(Tʳᶠ) #src
    write(io, " z_$(@sprintf("%.2e", Tʳᶠ[i]))") #src
end #src
write(io, " \n")  #src

for j = 1:length(Tᵢ) #src
    write(io, "$(@sprintf("%.2e", Tᵢ[j])) ") #src
    for i = 1:length(Tʳᶠ) #src
        write(io, "$(@sprintf("%.2e", M[j,i])) ") #src
    end #src
    write(io, " \n") #src
end #src
close(io) #src
#src #############################################################################

p = plot(xlabel="Tᵢ [s]", ylabel="zᶠ(Tʳᶠ, Tᵢ) [a.u.]")
for i = 1:length(TRF_scale)
    Mi = @view M[:,i]

    fit = curve_fit(standard_IR_model, Tᵢ, Mi, p0)
    param[i] = fit.param #src

    R₁[i] = fit.param[2]
    Minv[i] = fit.param[3] / fit.param[1] - 1
    residual[i] = norm(fit.resid) / norm(Mi)

    scatter!(p, Tᵢ, Mi, label=@sprintf("Tʳᶠ = %1.2es - data", Tʳᶠ[i]), color=i)
    plot!(p, Tᵢplot, standard_IR_model(Tᵢplot, fit.param), label=@sprintf("fit with R₁ = %.3f/s; MInv = %.3f", R₁[i], Minv[i]), color=i)
end
display(p) #!md
#md Main.HTMLPlot(p) #hide

#src #############################################################################
#src # export fitted curves
#src #############################################################################
Mp = [standard_IR_model(Tᵢplot, param[i]) for i=1:length(Tʳᶠ)] #src
io = open(expanduser(string("~/Documents/Paper/2021_MT_IDE/Figures/IR_monoExp_fit_BSA.txt")), "w") #src
write(io, "TI_s") #src
for i = 1:length(Tʳᶠ) #src
    write(io, " z_$(@sprintf("%.2e", Tʳᶠ[i]))") #src
end #src
write(io, " \n")  #src

for j = 1:length(Tᵢplot) #src
    write(io, "$(@sprintf("%.2e", Tᵢplot[j])) ") #src
    for i = 1:length(Tʳᶠ) #src
        write(io, "$(@sprintf("%.2e", Mp[i][j])) ") #src
    end #src
    write(io, " \n") #src
end #src
close(io) #src
#src #############################################################################

# Zooming into the early phase of the recovery curve reveals the poor fit quality, in particular for long ``T_\text{RF}``. This is also reflected by a substantially larger relative residual norm compared to the MnCl``_2`` sample:
mean(residual)

# In the paper, we highlight the estimated inversion efficiency and the relaxation rate of the dataset acquired with ``T_\text{RF}=22.8``μs
Minv[1]
#-
R₁[1] # 1/s
# and of the dataset acquired with ``T_\text{RF}=912``μs
Minv[end]
#-
R₁[end] # 1/s
# The mean value of all R₁ estimates is
mean(R₁) # 1/s

# and its standard deviation is substantially larger compared to the same fit of the MnCl``_2`` sample:
std(R₁) # 1/s

# ### Global IR Fit - Generalized Bloch Model
# In order to repeat the global fit that includes all ``T_\text{RF}`` values, we have to account for the spin dynamics in the semi-solid pool during the RF-pulse. First, we do this with the proposed generalized Bloch model:
function gBloch_IR_model(p, G, Tʳᶠ, TI, R2f)
    (m0, m0f_inv, m0s, R₁, T₂ˢ, Rx) = p
    m0f = 1 - m0s
    ω₁ = π ./ Tʳᶠ

    m0vec = [0, 0, m0f, m0s, 1]
    m_fun(p, t; idxs=nothing) = typeof(idxs) <: Number ? 0.0 : zeros(5)

    H = [-R₁-m0s*Rx     m0f*Rx R₁*m0f;
             m0s*Rx -R₁-m0f*Rx R₁*m0s;
              0          0         0 ]

    M = zeros(Float64, length(TI), length(Tʳᶠ))
    for i = 1:length(Tʳᶠ)
        param = (ω₁[i], 1, 0, m0s, R₁, R2f, Rx, R₁, T₂ˢ, G)
        prob = DDEProblem(apply_hamiltonian_gbloch!, m0vec, m_fun, (0.0, Tʳᶠ[i]), param)
        m = solve(prob)[end]

        for j = 1:length(TI)
            M[j,i] = m0 * (exp(H .* (TI[j] - Tʳᶠ[i] / 2)) * [m0f_inv * m[3],m[4],1])[1]
        end
    end
    return vec(M)
end;

# Here, we use assume a super-Lorentzian lineshape, whose Green's function is interpolated to speed up the fitting routine:
T₂ˢ_min = 5e-6 # s
G_superLorentzian = interpolate_greens_function(greens_superlorentzian, 0, maximum(Tʳᶠ)/T₂ˢ_min);

# The fit is initialized with `p0 = [m0, m0f_inv, m0_s, R₁, T2_s, Rx]` and we set some reasonable bounds to the fitted parameters:
p0   = [  1, 0.932,  0.1,   1, 10e-6, 50]
pmin = [  0, 0.100,   .0, 0.3,  1e-9, 10]
pmax = [Inf,   Inf,  1.0, Inf, 20e-6,1e3]

fit_gBloch = curve_fit((x, p) -> gBloch_IR_model(p, G_superLorentzian, Tʳᶠ, Tᵢ, 1/T₂star_BSA), [], vec(M), p0, lower=pmin, upper=pmax);

# Visually, the plot and the data align well:
p = plot(xlabel="Tᵢ [s]", ylabel="zᶠ(Tʳᶠ, Tᵢ) [a.u.]")
for i=1:length(Tʳᶠ)
    scatter!(p, Tᵢ, M[:,i], label=@sprintf("Tʳᶠ = %1.2es", Tʳᶠ[i]), color=i)
    plot!(p, Tᵢplot, gBloch_IR_model(fit_gBloch.param, G_superLorentzian, Tʳᶠ[i], Tᵢplot, 1/T₂star_BSA), label=@sprintf("Tʳᶠ = %1.2es", Tʳᶠ[i]), color=i)
end
display(p) #!md
#md Main.HTMLPlot(p) #hide

# which becomes particularly apparent when zooming into the beginning of the inversion recovery curves. Further, the relative residual norm is much smaller compared to the mono-exponential fit:
norm(fit_gBloch.resid) / norm(M)

# The estimated parameters are
m0 = fit_gBloch.param[1]
#-
Minv = fit_gBloch.param[2]
#-
m0s = fit_gBloch.param[3]
#-
R₁ = fit_gBloch.param[4] # 1/s
#-
T₂ˢ = 1e6fit_gBloch.param[5] # μs
#-
Rx = fit_gBloch.param[6] # 1/s
# with the uncertainties (in the same order)
stderror(fit_gBloch)

#src #############################################################################
#src # export fitted curves
#src #############################################################################
Mp = reshape(gBloch_IR_model(fit_gBloch.param, G_superLorentzian, Tʳᶠ, Tᵢplot, 1/T₂star_BSA), length(Tᵢplot), length(Tʳᶠ)) #src
io = open(expanduser(string("~/Documents/Paper/2021_MT_IDE/Figures/IR_gBloch_fit.txt")), "w") #src
write(io, "TI_s") #src
for i = 1:length(Tʳᶠ) #src
    write(io, " z_$(@sprintf("%.2e", Tʳᶠ[i]))") #src
end #src
write(io, " \n")  #src

for j = 1:length(Tᵢplot) #src
    write(io, "$(@sprintf("%.2e", Tᵢplot[j])) ") #src
    for i = 1:length(Tʳᶠ) #src
        write(io, "$(@sprintf("%.2e", Mp[j,i])) ") #src
    end #src
    write(io, " \n") #src
end #src
close(io) #src
#src #############################################################################

# ### Global IR Fit - Graham's Spectral Model
# For comparison, we repeat the same fit with [Graham's spectral model](http://dx.doi.org/10.1002/jmri.1880070520):
function Graham_IR_model(p, Tʳᶠ, TI, R2f)
    (m0, m0f_inv, m0s, R₁, T₂ˢ, Rx) = p
    m0f = 1 - m0s
    ω₁ = π ./ Tʳᶠ

    m0vec = [0, 0, m0f, m0s, 1]

    H = [-R₁-m0s*Rx     m0f*Rx R₁*m0f;
             m0s*Rx -R₁-m0f*Rx R₁*m0s;
              0          0         0 ]

    M = zeros(Float64, length(TI), length(Tʳᶠ))
    for i = 1:length(Tʳᶠ)
        param = (ω₁[i], 1, 0, Tʳᶠ[i], m0s, R₁, R2f, Rx, R₁, T₂ˢ)
        prob = ODEProblem(apply_hamiltonian_graham_superlorentzian!, m0vec, (0.0, Tʳᶠ[i]), param)
        m = solve(prob)[end]

        for j = 1:length(TI)
            M[j,i] = m0 * (exp(H .* (TI[j] - Tʳᶠ[i] / 2)) * [m0f_inv * m[3],m[4],1])[1]
        end
    end
    return vec(M)
end

fit_Graham = curve_fit((x, p) -> Graham_IR_model(p, Tʳᶠ, Tᵢ, 1/T₂star_BSA), [], vec(M), p0, lower=pmin, upper=pmax);

# Visually, the plot and the data align substantially worse:
p = plot(xlabel="Tᵢ [s]", ylabel="zᶠ(Tʳᶠ, Tᵢ) [a.u.]")
for i=1:length(Tʳᶠ)
    scatter!(p, Tᵢ, M[:,i], label=@sprintf("Tʳᶠ = %1.2es", Tʳᶠ[i]), color=i)
    plot!(p, Tᵢplot, Graham_IR_model(fit_Graham.param, Tʳᶠ[i], Tᵢplot, 1/T₂star_BSA), label=@sprintf("Tʳᶠ = %1.2es", Tʳᶠ[i]), color=i)
end
display(p) #!md
#md Main.HTMLPlot(p) #hide

# which becomes particularly apparent when zooming into the beginning of the inversion recovery curves. Further, the relative residual norm is much larger compared to the generalized Bloch fit:
norm(fit_Graham.resid) / norm(M)

# The estimated parameters are
m0 = fit_Graham.param[1]
#-
Minv = fit_Graham.param[2]
#-
m0s = fit_Graham.param[3]
#-
R₁ = fit_Graham.param[4] # 1/s
#-
T₂ˢ = 1e6fit_Graham.param[5] # μs
#-
Rx = fit_Graham.param[6] # 1/s
# with the uncertainties (in the same order)
stderror(fit_Graham)

#src #############################################################################
#src # export data
#src #############################################################################
#src export fitted curves
Mp = reshape(Graham_IR_model(fit_Graham.param, Tʳᶠ, Tᵢplot, 1/T₂star_BSA), length(Tᵢplot), length(Tʳᶠ)) #src
io = open(expanduser(string("~/Documents/Paper/2021_MT_IDE/Figures/IR_Graham_fit.txt")), "w") #src
write(io, "TI_s") #src
for i = 1:length(Tʳᶠ) #src
    write(io, " z_$(@sprintf("%.2e", Tʳᶠ[i]))") #src
end #src
write(io, " \n")  #src

for j = 1:length(Tᵢplot) #src
    write(io, "$(@sprintf("%.2e", Tᵢplot[j])) ") #src
    for i = 1:length(Tʳᶠ) #src
        write(io, "$(@sprintf("%.2e", Mp[j,i])) ") #src
    end #src
    write(io, " \n") #src
end #src
close(io) #src

# ### Global IR Fit - Sled's Model
# We also performed the fit with [Sled's model](http://dx.doi.org/10.1006/jmre.2000.2059):
function Sled_IR_model(p, G, Tʳᶠ, TI, R2f)
    (m0, m0f_inv, m0s, R₁, T₂ˢ, Rx) = p
    m0f = 1 - m0s
    ω₁ = π ./ Tʳᶠ

    m0vec = [0, 0, m0f, m0s, 1]

    H = [-R₁-m0s*Rx     m0f*Rx R₁*m0f;
             m0s*Rx -R₁-m0f*Rx R₁*m0s;
              0          0         0 ]

    M = zeros(Float64, length(TI), length(Tʳᶠ))
    for i = 1:length(Tʳᶠ)
        param = (ω₁[i], 1, 0, m0s, R₁, R2f, Rx, R₁, T₂ˢ, G)
        prob = ODEProblem(apply_hamiltonian_sled!, m0vec, (0.0, Tʳᶠ[i]), param)
        m = solve(prob)[end]

        for j = 1:length(TI)
            M[j,i] = m0 * (exp(H .* (TI[j] - Tʳᶠ[i] / 2)) * [m0f_inv * m[3],m[4],1])[1]
        end
    end
    return vec(M)
end

fit_Sled = curve_fit((x, p) -> Sled_IR_model(p, G_superLorentzian, Tʳᶠ, Tᵢ, 1/T₂star_BSA), [], vec(M), p0, lower=pmin, upper=pmax);

# Visually, the plot and the data do not align well either:
p = plot(xlabel="Tᵢ [s]", ylabel="zᶠ(Tʳᶠ, Tᵢ) [a.u.]")
for i=1:length(Tʳᶠ)
    scatter!(p, Tᵢ, M[:,i], label=@sprintf("Tʳᶠ = %1.2es", Tʳᶠ[i]), color=i)
    plot!(p, Tᵢplot, Sled_IR_model(fit_Sled.param, G_superLorentzian, Tʳᶠ[i], Tᵢplot, 1/T₂star_BSA), label=@sprintf("Tʳᶠ = %1.2es", Tʳᶠ[i]), color=i)
end
display(p) #!md
#md Main.HTMLPlot(p) #hide

# which becomes particularly apparent when zooming into the beginning of the inversion recovery curves. Further, the relative residual norm is also large compared to the generalized Bloch fit:
norm(fit_Sled.resid) / norm(M)

# The estimated parameters are
m0 = fit_Sled.param[1]
#-
Minv = fit_Sled.param[2]
#-
m0s = fit_Sled.param[3]
#-
R₁ = fit_Sled.param[4] # 1/s
#-
T₂ˢ = 1e6fit_Sled.param[5] # μs
#-
Rx = fit_Sled.param[6] # 1/s
# with the uncertainties (in the same order)
stderror(fit_Sled)

#src #############################################################################
#src # export data
#src #############################################################################
#src export fitted curves
Mp = reshape(Sled_IR_model(fit_Sled.param, G_superLorentzian, Tʳᶠ, Tᵢplot, 1/T₂star_BSA), length(Tᵢplot), length(Tʳᶠ)) #src
io = open(expanduser(string("~/Documents/Paper/2021_MT_IDE/Figures/IR_Sled_fit.txt")), "w") #src
write(io, "TI_s") #src
for i = 1:length(Tʳᶠ) #src
    write(io, " z_$(@sprintf("%.2e", Tʳᶠ[i]))") #src
end #src
write(io, " \n")  #src

for j = 1:length(Tᵢplot) #src
    write(io, "$(@sprintf("%.2e", Tᵢplot[j])) ") #src
    for i = 1:length(Tʳᶠ) #src
        write(io, "$(@sprintf("%.2e", Mp[j,i])) ") #src
    end #src
    write(io, " \n") #src
end #src
close(io) #src

# ### Analysis of the Residuals
# In order to visualize how well the three models align with the data at different ``T_\text{RF}``, we calculate the ``\ell_2``-norm of the residuals after subtracting the modeled from the measured signal and normalize it by the ``\ell_2``-norm of the signal:

resid_gBlo = similar(Tʳᶠ)
resid_Sled = similar(Tʳᶠ)
resid_Grah = similar(Tʳᶠ)
for i=1:length(Tʳᶠ)
    resid_gBlo[i] = norm(gBloch_IR_model(fit_gBloch.param, G_superLorentzian, Tʳᶠ[i], Tᵢ, 1/T₂star_BSA) .- M[:,i]) / norm(M[:,i])
    resid_Grah[i] = norm(Graham_IR_model(fit_Graham.param, Tʳᶠ[i], Tᵢ, 1/T₂star_BSA)                    .- M[:,i]) / norm(M[:,i])
    resid_Sled[i] = norm(Sled_IR_model(  fit_Sled.param,   G_superLorentzian, Tʳᶠ[i], Tᵢ, 1/T₂star_BSA) .- M[:,i]) / norm(M[:,i])
end

p = plot(xlabel="Tʳᶠ [s]", ylabel="relative residual")
scatter!(p, Tʳᶠ, resid_gBlo, label="generalized Bloch model")
scatter!(p, Tʳᶠ, resid_Grah, label="Graham's spectral model")
scatter!(p, Tʳᶠ, resid_Sled, label="Sled's model")
#md Main.HTMLPlot(p) #hide

#src #############################################################################
#src # export data
#src #############################################################################
io = open(expanduser(string("~/Documents/Paper/2021_MT_IDE/Figures/IR_residuum_each_fit.txt")), "w") #src
write(io, "TRF_ms Graham_percent Sled_percent gBloch_percent \n")  #src

for i = 1:length(Tʳᶠ) #src
    write(io, "$(@sprintf("%.2e", 1e3 * Tʳᶠ[i])) ") #src
    write(io, "$(@sprintf("%.2e", 1e2 * resid_Grah[i])) ") #src
    write(io, "$(@sprintf("%.2e", 1e2 * resid_Sled[i])) ") #src
    write(io, "$(@sprintf("%.2e", 1e2 * resid_gBlo[i])) ") #src
    write(io, " \n") #src
end #src
close(io) #src

# This analysis examines the residuals from the actual fits, i.e. it uses the biophysical parameters of respective fit to model the signal. The disadvantage of this approach is that residuals at long ``T_\text{RF}`` are negatively affected by the poor fits of Graham's and Sled's models at short ``T_\text{RF}``. This problem is overcome by subtracting the measured signal from signal that is simulated with the biophysical parameters that were estimated by fitting the generalized Bloch model:

for i=1:length(Tʳᶠ)
    resid_gBlo[i] = norm(gBloch_IR_model(fit_gBloch.param, G_superLorentzian, Tʳᶠ[i], Tᵢ, 1/T₂star_BSA) .- M[:,i]) / norm(M[:,i])
    resid_Grah[i] = norm(Graham_IR_model(fit_gBloch.param, Tʳᶠ[i], Tᵢ, 1/T₂star_BSA)                    .- M[:,i]) / norm(M[:,i])
    resid_Sled[i] = norm(Sled_IR_model(  fit_gBloch.param, G_superLorentzian, Tʳᶠ[i], Tᵢ, 1/T₂star_BSA) .- M[:,i]) / norm(M[:,i])
end

p = plot(xlabel="Tʳᶠ [s]", ylabel="relative residual")
scatter!(p, Tʳᶠ, resid_gBlo, label="generalized Bloch model")
scatter!(p, Tʳᶠ, resid_Grah, label="Graham's spectral model")
scatter!(p, Tʳᶠ, resid_Sled, label="Sled's model")
#md Main.HTMLPlot(p) #hide

# One can observe reduced residuals for Graham's and Sled's models for long ``T_\text{RF}`` as a trade off for larger residuals at short ``T_\text{RF}``. Yet, the residuals at long ``T_\text{RF}`` are still substantially larger compared to ones of the generalized Bloch model.

#src #############################################################################
#src # export data
#src #############################################################################
io = open(expanduser(string("~/Documents/Paper/2021_MT_IDE/Figures/IR_residuum_gBloch_est.txt")), "w") #src
write(io, "TRF_ms Graham_percent Sled_percent gBloch_percent \n")  #src

for i = 1:length(Tʳᶠ) #src
    write(io, "$(@sprintf("%.2e", 1e3 * Tʳᶠ[i])) ") #src
    write(io, "$(@sprintf("%.2e", 1e2 * resid_Grah[i])) ") #src
    write(io, "$(@sprintf("%.2e", 1e2 * resid_Sled[i])) ") #src
    write(io, "$(@sprintf("%.2e", 1e2 * resid_gBlo[i])) ") #src
    write(io, " \n") #src
end #src
close(io) #src


struct TrapezoidalPoints{T <: AbstractFloat} <: AbstractVector{T}
    data::LinRange{T}
    function TrapezoidalPoints{T}(n) where {T}
        if n <= 1
            return new{T}(LinRange{T}(0,0,n))
        else
            return new{T}(LinRange{T}(-1,1,n))
        end
    end
end
TrapezoidalPoints(n) = TrapezoidalPoints{Float64}(n)
Base.size(x::TrapezoidalPoints) = size(x.data)
Base.getindex(x::TrapezoidalPoints, i::Int) = x.data[i]


include("_common.jl")

using StatsPlots, OnlineStats

calib_start = 1826  # 1826
outflows = Dict()
climate = FULL_CLIMATE
calib_obs = CALIB_OBS[calib_start+1:end]

baseline_path = joinpath(FIG_PATH, "calib_baseline")

for app in APPROACHES
    display_name = replace(app, "_"=>" ")
    sn, n_id = setup_network(joinpath("baselines", "cotter_baseline_IHACRES_$(app).yml"))
    node = sn[n_id]
    run_basin!(sn, climate)

    calib_outflow = node.outflow[calib_start+1:CALIB_LN]
    gw_store = node.gw_store[calib_start+1:CALIB_LN]

    cmd = node.storage[calib_start+1:CALIB_LN]
    log_cmd = log.(cmd)
    max_cmd = maximum(log_cmd)
    quantiles = [0.2, 0.80]

    normal_dist = fit(Normal, log_cmd)
    μ = mean(normal_dist)
    σ = std(normal_dist)
    d_min = minimum(log_cmd)
    d_max = maximum(log_cmd)
    trunc_dist = truncated(Normal(μ, σ), d_min, d_max)

    thresholds = quantile(log_cmd, quantiles)
    @info "Log CMD thresholds: [0.1, 0.9 / min / max]" thresholds d_min d_max
    histogram(log_cmd, alpha=0.4, normalize=true,
              title="Distribution of Log CMD values\n($(display_name) calibrated baseline)",
              legend=:topleft,
              label=false,
              xlabel="Log CMD",
              ylabel="Density")
    StatsPlots.plot!(trunc_dist, linewidth=2, label="Fitted normal distribution")
    vline!(thresholds; color="orange", linewidth=1.5, label="State thresholds")
    for (thres, annot) in zip(thresholds, quantiles)
        annotate!(thres, 0.0015, text("q$annot", 8))
    end

    savefig(joinpath(baseline_path, "baseline_cmd_dist_$(app).png"))

    # Save performance over calibration period
    rmse = round(Streamfall.RMSE(calib_obs, calib_outflow), digits=2)
    mkge = round(Streamfall.mKGE(calib_obs, calib_outflow), digits=2)
    nse = round(Streamfall.NSE(calib_obs, calib_outflow), digits=2)
    nnse = round(Streamfall.NNSE(calib_obs, calib_outflow), digits=2)
    nmkge = round(Streamfall.NmKGE(calib_obs, calib_outflow), digits=2)
    mean_nmkge = round(Streamfall.mean_NmKGE(calib_obs, calib_outflow), digits=2)
    @info "$app calibration: "
    report_metrics(calib_obs, calib_outflow)

    plot(CALIB_DATES[calib_start+1:end], calib_obs, label="Historic", title="$app Calibration\n(RMSE: $rmse; NSE: $nse; KGE': $mkge)", legend=:topright)
    plot!(CALIB_DATES[calib_start+1:end], calib_outflow, label="Simulated", alpha=0.7, ylabel="Stream flow [ML]")
    savefig(joinpath(baseline_path, "baseline_calib_perf_$(app).png"))

    qqplot(calib_obs, calib_outflow, title="Calibration Period\n$display_name")
    xlabel!("Historic")
    ylabel!("Simulated")
    savefig(joinpath(baseline_path, "baseline_calib_perf_qq_$(display_name).png"))

    @info "CMD" app minimum(node.storage[2:end]) maximum(node.storage[2:end])
    @info "GW Store" app minimum(node.gw_store[2:end]) maximum(node.gw_store[2:end])
    @info "Quick" app minimum(node.quick_store[2:end]) maximum(node.quick_store[2:end])
    @info "Slow" app minimum(node.slow_store[2:end]) maximum(node.slow_store[2:end])

    # Save performance over validation period
    # valid_cmd = node.storage[CALIB_LN+1:end]
    valid_outflow = node.outflow[CALIB_LN+1:end]
    outflows[app] = valid_outflow

    rmse = round(Streamfall.RMSE(VALID_OBS, valid_outflow), digits=2)
    mkge = round(Streamfall.mKGE(VALID_OBS, valid_outflow), digits=2)
    nse = round(Streamfall.NSE(VALID_OBS, valid_outflow), digits=2)
    nnse = round(Streamfall.NNSE(VALID_OBS, valid_outflow), digits=2)
    nmkge = round(Streamfall.NmKGE(VALID_OBS, valid_outflow), digits=2)
    mean_nmkge = round(Streamfall.mean_NmKGE(VALID_OBS, valid_outflow), digits=2)
    @info "$app validation: "
    report_metrics(VALID_OBS, valid_outflow)

    plot(VALID_DATES, VALID_OBS, label="Historic", title="$display_name Validation\n(RMSE: $rmse; NSE: $nse; KGE': $mkge)", legend=:topright)
    plot!(VALID_DATES, valid_outflow, label="Simulated", alpha=0.7, ylabel="Stream flow [ML]")
    savefig(joinpath(baseline_path, "baseline_valid_perf_$(app).png"))

    qqplot(VALID_OBS, valid_outflow, title="Validation Period\n$display_name")
    xlabel!("Historic")
    ylabel!("Simulated")
    savefig(joinpath(baseline_path, "baseline_valid_perf_qq_$(app).png"))

    full_dates = vcat(CALIB_DATES, VALID_DATES)
    plot(full_dates[calib_start:CALIB_LN], CALIB_OBS[calib_start:end], label="Calibration", title="$display_name C&V", legend=:topright)
    plot!(full_dates[CALIB_LN+1:end], VALID_OBS, label="Validation", color="green")
    plot!(full_dates[calib_start:end], node.outflow[calib_start:end], label="Simulated", linestyle=:dashdot, color="red", alpha=0.7, ylabel="Stream flow [ML]")
    savefig(joinpath(baseline_path, "baseline_full_perf_$(app).png"))
end

# Ensemble results
ensemble_outflow = mean.(outflows["NNSE"] .+ outflows["mean_NmKGE"])
rmse = round(Streamfall.RMSE(VALID_OBS, ensemble_outflow), digits=2)
mkge = round(Streamfall.mKGE(VALID_OBS, ensemble_outflow), digits=2)
nse = round(Streamfall.NSE(VALID_OBS, ensemble_outflow), digits=2)

plot(VALID_DATES, VALID_OBS, label="Historic", title="Baseline Ensemble\n(RMSE: $rmse; NSE: $nse; KGE': $mkge)", legend=:topright)
plot!(VALID_DATES, ensemble_outflow, label="Ensemble", alpha=0.8)
savefig(joinpath(FIG_PATH, "baseline_valid_perf_ensemble_NNSE_meanNmKGE.png"))

# keeping the below as it plots out simulation for each bin
# climate = Climate(training, "_rain", "_temp")

# # Set up stream network
# network = YAML.load_file(joinpath(DATA_PATH, "cotter_baseline_calibrated_NNSE.yml"))
# nnse_sn = create_network("Calibrated NNSE", network)

# run_basin!(nnse_sn, climate)

# node = nnse_sn[1]
# nnse_outflow = node.outflow
# nnse_cmd = node.storage[2:end]  # skip the first entry as it is the initial condition

# normal_dist = fit(Normal, nnse_cmd)
# trunc_dist = truncated(Normal(mean(normal_dist), std(normal_dist)), 0.0, maximum(nnse_cmd))

# # save to file
# dist_data = Dict(
#     "model" => ["Baseline_NNSE", "Baseline_NmKGE", "Baseline_RMSE"],
#     "mean" => [mean(normal_dist), ],
#     "std" => [std(normal_dist), ],
#     "max" => [maximum(nnse_cmd), ],
# )

# nnse_bin_bounds = quantile.(trunc_dist, [0.5-0.25, 0.5+0.25, 1.0])
# histogram(nnse_cmd, alpha=0.4, normalize=true,
#           title="Distribution of CMD values\n(NNSE calibrated baseline)",
#           legend=:left,
#           label="CMD",
#           xlabel="Catchment Moisture Deficit",
#           ylabel="Probability")
# StatsPlots.plot!(trunc_dist, linewidth=2, label="Fitted normal distribution")
# vline!(nnse_bin_bounds; color="orange", linewidth=1.5, label="Ensemble bounds")

# savefig(joinpath(FIG_PATH, "NNSE_calibrated_dist.png"))

# ######

# # Set up stream network
# nmkge_network = YAML.load_file(joinpath(DATA_PATH, "cotter_baseline_calibrated_NmKGE.yml"))
# nmkge_sn = create_network("Calibrated NmKGE", nmkge_network)

# run_basin!(nmkge_sn, climate)

# node = nmkge_sn[1]
# nmkge_outflow = node.outflow
# nmkge_cmd = node.storage[2:end]

# normal_dist = fit(Normal, nmkge_cmd)
# trunc_dist = truncated(Normal(mean(normal_dist), std(normal_dist)), 0.0, maximum(nmkge_cmd))

# # Save data
# append!(dist_data["mean"], mean(normal_dist))
# append!(dist_data["std"], std(normal_dist))
# append!(dist_data["max"], maximum(nmkge_cmd))

# nmkge_bin_bounds = quantile.(trunc_dist, [0.5-0.25, 0.5+0.25, 1.0])
# histogram(nmkge_cmd, alpha=0.4, normalize=true,
#           title="Distribution of CMD values\n(NmKGE calibrated baseline)",
#           legend=:left,
#           label="CMD",
#           xlabel="Catchment Moisture Deficit",
#           ylabel="Probability")
# StatsPlots.plot!(trunc_dist, linewidth=2, label="Fitted normal distribution")
# vline!(nmkge_bin_bounds; color="orange", linewidth=1.5, label="Ensemble bounds")

# savefig(joinpath(FIG_PATH, "NmKGE_calibrated_dist.png"))

# ######

# # Set up stream network
# rmse_network = YAML.load_file(joinpath(DATA_PATH, "cotter_baseline_calibrated_RMSE.yml"))
# rmse_sn = create_network("Calibrated RMSE", rmse_network)

# run_basin!(rmse_sn, climate)

# node = rmse_sn[1]
# rmse_outflow = node.outflow
# rmse_cmd = node.storage[2:end]

# normal_dist = fit(Normal, rmse_cmd)
# trunc_dist = truncated(Normal(mean(normal_dist), std(normal_dist)), 0.0, maximum(rmse_cmd))

# # Save data
# append!(dist_data["mean"], mean(normal_dist))
# append!(dist_data["std"], std(normal_dist))
# append!(dist_data["max"], maximum(rmse_cmd))
# CSV.write(joinpath(DATA_PATH, "dist_data.csv"), DataFrame(dist_data))

# nmkge_bin_bounds = quantile.(trunc_dist, [0.5-0.25, 0.5+0.25, 1.0])
# histogram(rmse_cmd, alpha=0.4, normalize=true,
#           title="Distribution of CMD values\n(RMSE calibrated baseline)",
#           legend=:left,
#           label="CMD",
#           xlabel="Catchment Moisture Deficit",
#           ylabel="Probability")
# StatsPlots.plot!(trunc_dist, linewidth=2, label="Fitted normal distribution")
# vline!(nmkge_bin_bounds; color="orange", linewidth=1.5, label="Ensemble bounds")

# savefig(joinpath(FIG_PATH, "RMSE_calibrated_dist.png"))

# #######

# # Performance of baseline and ensemble of baselines
# hist_flow = training_flow["410730"]
# mean_outflow = mean([nnse_outflow, nmkge_outflow], dims=1)[1]
# @info "Performance of baseline and ensemble of baselines"
# @info "NNSE Calibrated" Streamfall.RMSE(hist_flow, nnse_outflow)
# @info "NmKGE Calibrated" Streamfall.RMSE(hist_flow, nmkge_outflow)
# @info "RMSE Calibrated" Streamfall.RMSE(hist_flow, rmse_outflow)
# @info "Mean of Ensemble" Streamfall.RMSE(hist_flow, mean_outflow)

# # Performance within each bin
# nnse_bin_membership = bin_membership(nnse_cmd, nnse_bin_bounds)
# @info "Performance within bins (NNSE)"
# for idx in 1:length(nnse_bin_bounds)+1
#     bin_member = findall(x -> x == idx, nnse_bin_membership)
#     @info "Bin $idx ($(length(bin_member)))"
#     @info "NNSE" Streamfall.RMSE(hist_flow[bin_member], nnse_outflow[bin_member])
#     @info "Ensemble" Streamfall.RMSE(hist_flow[bin_member], mean_outflow[bin_member])
# end

# ######
# println("------------")

# nmkge_bin_membership = bin_membership(nmkge_cmd, nmkge_bin_bounds)
# @info "Performance within bins (NmKGE)"
# for idx in 1:length(nmkge_bin_bounds)+1
#     bin_member = findall(x -> x == idx, nmkge_bin_membership)
#     @info "Bin $idx ($(length(bin_member)))"
#     if length(bin_member) == 0
#         @info "No members"
#         continue
#     end
#     @info "NmKGE" Streamfall.RMSE(hist_flow[bin_member], nmkge_outflow[bin_member])
#     @info "Ensemble" Streamfall.RMSE(hist_flow[bin_member], mean_outflow[bin_member])
# end

# plot(hist_flow, title="Baseline results", label="Observed",
#      xlabel="Time step", ylabel="Streamflow [ML]")
# num_bounds = length(nnse_bin_bounds)+1
# for idx in 1:num_bounds
#     bin_member = findall(x -> x == idx, nnse_bin_membership)
#     if length(bin_member) == 0
#         continue
#     end
#     scatter!(bin_member, nnse_outflow[bin_member],
#              markersize=2, label="Bin $idx",
#              markerstrokewidth=0,
#              alpha=0.6;
#              palette=cgrad(:matter, num_bounds, categorical = true))
# end

# savefig(joinpath(FIG_PATH, "NNSE_baseline_perf.png"))

# plot(hist_flow, title="Baseline results", label="Observed",
#      xlabel="Time step", ylabel="Streamflow [ML]")
# for idx in 1:num_bounds
#     bin_member = findall(x -> x == idx, nmkge_bin_membership)
#     if length(bin_member) == 0
#         continue
#     end
#     scatter!(bin_member, nmkge_outflow[bin_member],
#              markersize=2, label="Bin $idx",
#              markerstrokewidth=0,
#              alpha=0.6;
#              palette=cgrad(:matter, num_bounds, categorical = true))
# end

# savefig(joinpath(FIG_PATH, "NmKGE_baseline_perf.png"))

# v_flow = validation_flow["410730"]
# reset!(nnse_sn)
# v_climate = Climate(validation, "_rain", "_temp")
# run_basin!(nnse_sn, v_climate)

# plot(v_flow, label="Observed", alpha=0.5, title="Validation Period")
# nnse_outflow = nnse_sn[1].outflow
# rmse = round(Streamfall.RMSE(v_flow, nnse_outflow), digits=4)
# plot!(nnse_outflow, label="NNSE calibrated: RMSE = $rmse", linestyle=:dash, alpha=0.9)

# ###
# reset!(nmkge_sn)
# run_basin!(nmkge_sn, v_climate)
# nmkge_outflow = nmkge_sn[1].outflow
# rmse = round(Streamfall.RMSE(v_flow, nmkge_outflow), digits=4)
# plot!(nmkge_outflow, label="NmKGE calibrated: RMSE = $rmse", linestyle=:dashdotdot, alpha=0.8)

# ####
# reset!(rmse_sn)
# run_basin!(rmse_sn, v_climate)
# rmse_outflow = rmse_sn[1].outflow
# rmse = round(Streamfall.RMSE(v_flow, rmse_outflow), digits=4)
# plot!(rmse_outflow, label="RMSE calibrated: RMSE = $rmse", linestyle=:dashdotdot, alpha=0.7)

# # mean_outflow = mean([nnse_outflow, nmkge_outflow], dims=1)[1]
# # rmse = round(Streamfall.RMSE(v_flow, mean_outflow), digits=4)
# # plot!(mean_outflow, label="Ensemble mean: RMSE = $rmse")

# savefig(joinpath(FIG_PATH, "calibrated_performance.png"))


"""
Nlayer_cylinder_DCS

Provides default parameters to use to compute the autocorrelation function g2 in a N-layered cylinder.
The arguments defined with a Vector should be the same length and same type.

# Keyword Arguments
- `ρ`: source-detector separation (cm⁻¹)
- `μa::Vector{T}`: absorption coefficient (cm⁻¹)
- `μsp::Vector{T}`: reduced scattering coefficient (cm⁻¹)
- `n_med::Vector{T}`: medium's index of refraction
- `n_ext`: external medium's index of refraction (air or detector)
- `l`: the thicknesses of each layer (cm)
- `z`: the z-depth orthogonal from the boundary (cm)
- `a`: the radius of the cylinder (cm)
- `β`: constant in Siegert relation dependent on collection optics
- `BFi::Vector{T}`: Blood flow index ~αDb (cm²/s)
- `λ`: wavelength (nm)
- `bessels::AbstractVector{T}`: roots of J0

# Examples
```
julia> data = Nlayer_cylinder_DCS() # return default parameters
julia> data = Nlayer_cylinder_DCS(ρ = 1.5) # return ρ = 1.5 with the rest of the parameters given by defaults

# we can now define our correlation times τ
julia> τ = 10 .^(range(-10,stop=0,length=250))
julia> g2_DA_Nlay_cylinder_CW(τ, data)
```
"""
@with_kw struct Nlayer_cylinder_DCS{T <: Real}
    μsp::Vector{T} = [10.0, 10.0]                       # reduced scattering coefficient (1/cm)
    μa::Vector{T} = [0.1, 0.1]                          # absorption coefficient (1/cm)
    n_ext::T = 1.0                                      # surrounding index of refraction
    n_med::Vector{T} = [1.0, 1.0]                       # layers index of refraction

    l::Vector{T} = [0.5, 10.0]                          # length of cylinder layers (cm)
    ρ::Union{T, AbstractVector{T}} = 1.0                # source-detector separation (cm)
    a::T = 20.0                                         # radius of cylinder (cm)
    z::T = 0.0                                          # detector depth (cm)

    β::T = 1.0                                          # constant in Siegert relation dependent on collection optics
    BFi::Vector{T} = [2.0e-8, 2.0e-8]                   # Blood flow index ~αDb (cm²/s)
    λ::T = 750.0                                        # wavelength (nm)

    bessels::AbstractVector{T} = besselroots[1:1000]    # roots of J0

    N_quad::Int = 50                                    # number of nodes in gauss-legendre integration
    N_laplace::Int = 12                                 # number of laplace evaluations in time-domain inversion
end

"""
    g1_DA_Nlay_cylinder_CW(τ::AbstractVector, ρ, μa, μsp; BFi = [2e-8, 2e-8], n_ext = 1.0, n_med = [1.0, 1.0], l = [1.0, 10.0], a = 20.0, z = 0.0, λ = 700.0, bessels = besselroots[1:500])

Compute the electric field autocorrelation function function g1 in a layered cylinder (CW).

# Arguments
- `τ`: time lags (s)
- `ρ`: the source detector separation (cm⁻¹)
- `μa`: absorption coefficient (cm⁻¹)
- `μsp`: reduced scattering coefficient (cm⁻¹)

# Keyword arguments
- `BFi`: Blood flow index ~αDb (cm²/s)
- `β`: constant in Siegert relation dependent on collection optics
- `n_ext`: the boundary's index of refraction (air or detector)
- `n_med`: the sample medium's index of refraction
- `l`: length of cylinder layers (cm)
- `a`: cylinder radius (cm)
- `z`: the z-depth orthogonal from the boundary (cm)
- `λ`: wavelength (nm)
- `bessels`: roots of J0 (bessel function of first kind order zero)
"""
function g1_DA_Nlay_cylinder_CW(τ::AbstractVector, ρ, μa, μsp; BFi = [2e-8, 2e-8], n_ext = 1.0, n_med = [1.0, 1.0], l = [1.0, 10.0], a = 20.0, z = 0.0, λ = 700.0, bessels = besselroots[1:500])
    g1 = similar(τ)
    μa_dynamic = similar(μa)

    k0 = 2 * π * n_med / (λ * 1e-7) # this should be in cm
    tmp = @. 2 * μsp * BFi * k0^2

    G0 = fluence_DA_Nlay_cylinder_CW(ρ, μa, μsp, n_ext, n_med, l, a, z, bessels)

    Threads.@threads for ind in eachindex(τ)
        μa_dynamic = μa + tmp * τ[ind]
        G1 = fluence_DA_Nlay_cylinder_CW(ρ, μa_dynamic, μsp, n_ext, n_med, l, a, z, bessels)
        g1[ind] = G1 / G0
    end

    return g1
end

"""
    g1_DA_Nlay_cylinder_TD(τ::AbstractVector, t::AbstractFloat, ρ, μa, μsp; BFi = [2e-8, 2e-8], n_ext = 1.0, n_med = [1.0, 1.0], l = [1.0, 10.0], a = 20.0, z = 0.0, λ = 700.0, bessels = besselroots[1:1000], N_laplace = 8)
    g1_DA_Nlay_cylinder_TD(τ::AbstractVector, t::AbstractVector, ρ, μa, μsp; BFi = [2e-8, 2e-8], n_ext = 1.0, n_med = [1.0, 1.0], l = [1.0, 10.0], a = 20.0, z = 0.0, λ = 700.0, bessels = besselroots[1:1000], N_laplace = 8, N_quad = 100)

Compute the time-domain electric field autocorrelation function function g1 in a layered cylinder.

If t is a AbstractVector then g1(tau) will be calculated by ∫P(t)g1(tau, t)dt with integration bounds given by t[1] and t[end].
This is needed for time-gated applications.

# Arguments
- `τ`: time lags (s)
- `ρ`: the source detector separation (cm⁻¹)
- `μa`: absorption coefficient (cm⁻¹)
- `μsp`: reduced scattering coefficient (cm⁻¹)

# Keyword arguments
- `BFi`: Blood flow index ~αDb (cm²/s)
- `β`: constant in Siegert relation dependent on collection optics
- `n_ext`: the boundary's index of refraction (air or detector)
- `n_med`: the sample medium's index of refraction
- `l`: length of cylinder layers (cm)
- `a`: cylinder radius (cm)
- `z`: the z-depth orthogonal from the boundary (cm)
- `λ`: wavelength (nm)
- `bessels`: roots of J0 (bessel function of first kind order zero)
- `N_quad`: number of nodes in gauss-legendre integration
- `N_laplace`: number of laplace evaluations in inversion to time-domain
"""
function g1_DA_Nlay_cylinder_TD(τ::AbstractVector, t::AbstractFloat, ρ, μa, μsp; BFi = [2e-8, 2e-8], n_ext = 1.0, n_med = [1.0, 1.0], l = [1.0, 10.0], a = 20.0, z = 0.0, λ = 700.0, bessels = besselroots[1:1000], N_laplace = 8)
    g1 = similar(τ)
    μa_dynamic = similar(μa)

    k0 = 2 * π * n_med / (λ * 1e-7) # this should be in cm
    tmp = @. 2 * μsp * BFi * k0^2

    G0 = fluence_DA_Nlay_cylinder_TD(t, ρ, μa, μsp, n_ext, n_med, l, a, z, bessels, N = N_laplace)

    for ind in eachindex(τ)
        μa_dynamic = μa + tmp * τ[ind]
        G1 = fluence_DA_Nlay_cylinder_TD(t, ρ, μa_dynamic, μsp, n_ext, n_med, l, a, z, bessels, N = N_laplace)
        g1[ind] = G1 / G0
    end

    return g1
end

function g1_DA_Nlay_cylinder_TD(τ::AbstractVector, t::AbstractVector, ρ, μa, μsp; BFi = [2e-8, 2e-8], n_ext = 1.0, n_med = [1.0, 1.0], l = [1.0, 10.0], a = 20.0, z = 0.0, λ = 700.0, bessels = besselroots[1:1000], N_laplace = 8, N_quad = 100)
    g1 = similar(τ)
    μa_dynamic = similar(μa)

    k0 = 2 * π * n_med / (λ * 1e-7) # this should be in cm
    tmp = @. 2 * μsp * BFi * k0^2
    x, w = gausslegendre(N_quad)

    tpsf_norm = integrate_compute_y_first(t -> fluence_DA_Nlay_cylinder_TD(t, ρ, μa, μsp, n_ext, n_med, l, a, z, bessels, N = N_laplace), x, w, t[1], t[2])

    for ind in eachindex(τ)
        μa_dynamic = μa + tmp * τ[ind]
        g1[ind] = integrate_compute_y_first(t -> (fluence_DA_Nlay_cylinder_TD(t, ρ, μa_dynamic, μsp, n_ext, n_med, l, a, z, bessels, N = N_laplace) / tpsf_norm), x, w, t[1], t[2])
    end

    return g1
end

function g2_DA_Nlay_cylinder_CW(τ::AbstractVector, ρ, μa, μsp; BFi = [2e-8, 2e-8], n_ext = 1.0, n_med = [1.0, 1.0], l = [1.0, 10.0], a = 20.0, z = 0.0, λ = 700.0, bessels = besselroots[1:500], β = 1.0)
    g1 = g1_DA_Nlay_cylinder_CW(τ, ρ, μa, μsp, BFi = BFi, n_ext = n_ext, n_med = n_med, l = l, a = a, z = z, λ = λ, bessels = bessels)
    for ind in eachindex(τ)
        g1[ind] = 1 + β * abs(g1[ind])^2
    end
    return g1
end
function g2_DA_Nlay_cylinder_CW(τ::AbstractVector, data::Nlayer_cylinder_DCS)
    return g2_DA_Nlay_cylinder_CW(τ, data.ρ, data.μa, data.μsp, BFi = data.BFi, n_ext = data.n_ext, n_med = data.n_med, l = data.l, a = data.a, λ = data.λ, bessels = data.bessels, β = data.β)
end

function g2_DA_Nlay_cylinder_TD(τ::AbstractVector, t::Union{AbstractVector, AbstractFloat}, ρ, μa, μsp; BFi = [2e-8, 2e-8], n_ext = 1.0, n_med = [1.0, 1.0], l = [1.0, 10.0], a = 20.0, z = 0.0, λ = 700.0, bessels = besselroots[1:500], N_laplace = 12, N_quad = 50, β = 1.0)
    if t isa AbstractVector
        g1 = g1_DA_Nlay_cylinder_TD(τ, t, ρ, μa, μsp, BFi = BFi, n_ext = n_ext, n_med = n_med, l = l, a = a, z = z, λ = λ, bessels = bessels, N_laplace = N_laplace, N_quad = N_quad)
    elseif t isa AbstractFloat
        g1 = g1_DA_Nlay_cylinder_TD(τ, t, ρ, μa, μsp, BFi = BFi, n_ext = n_ext, n_med = n_med, l = l, a = a, z = z, λ = λ, bessels = bessels, N_laplace = N_laplace)
    end

    for ind in eachindex(τ)
        g1[ind] = 1 + β * abs(g1[ind])^2
    end
    return g1
end
function g2_DA_Nlay_cylinder_TD(τ::AbstractVector, t, data::Nlayer_cylinder_DCS)
    return g2_DA_Nlay_cylinder_TD(τ, t, data.ρ, data.μa, data.μsp, BFi = data.BFi, n_ext = data.n_ext, n_med = data.n_med, l = data.l, a = data.a, λ = data.λ, bessels = data.bessels, N_laplace = data.N_laplace, N_quad = data.N_quad, β = data.β)
end


#############################################################################
# lovasz_extension.jl
# Handles the Lovasz extensions of set functions.
#############################################################################

export lovasz, LovaszExtAbsAtom
export sign, monotonicity, curvature, evaluate
export ConvexLovaszAbs

mutable struct LovaszExtAbsAtom <: AbstractExpr
  head::Symbol
  id_hash::UInt64
  children::Tuple{AbstractExpr, Variable}
  size::Tuple{Int, Int}
  variable::Variable
  func::SubmodFunc

  function LovaszExtAbsAtom(F::SubmodFunc, x::AbsAtom)
    if length(F.setvariables) == 1
      if length(x.children) == 1 && isa(x.children[1], Variable)
        if evaluate(F, [])[1] != 0
          error("A combinatorial function should be 0 at the empty set to derive its Lovasz extension.")
        else
          var = get_v(x)
          if length(var) != 1
            error("Functions with other than one variable are not supported.")
          else
            if var[1].size[1] == F.setvariables[1].cardinality
              children = (F, var[1])
              return new(:lovaszabs, hash(children), children, (1, 1), var[1], F)
            else
              error("The size of the continuous variable should be the same as the baseset of the combinatorial variable of the combinatorial function.")
            end
          end
        end
      else
        error("Only able to define Lovasz extesions on absolute values of a variable.")
      end
    else
      error("A combinatorial function should be sigle-variant to obtain its Lovasz extension.")
    end
  end
end

lovasz(F::SubmodFunc, x::AbsAtom) = LovaszExtAbsAtom(F, x)

function sign(F::LovaszExtAbsAtom)
  return sign(F.children[1])
end

function monotonicity(F::LovaszExtAbsAtom)
  return monotonicity(F.children[1])
end

function curvature(F::LovaszExtAbsAtom)
  modd = modularity(F.children[1])
  if modd == SubModularity()
    return ConvexVexity()
  elseif modd == SuperModularity()
    return ConcaveVexity()
  elseif modd == Modularity()
    return AffineVexity()
  elseif modd == ConstModularity()
    return ConstVexity()
  end
end

vexity(F::LovaszExtAbsAtom) = curvature(F)

function evaluate(f::LovaszExtAbsAtom)
  n = f.children[2].size[1]
  y = abs.(f.children[2].value[:])
  i = sortperm(y, rev = true)
  V = sort(f.children[1].setvariables[1].baseset)
  storage = zeros(n + 1)
  x = zeros(n)
  for ii = 2:n+1
    storage[ii] = evaluate(f.children[1], V[i[1:ii-1]])[1]
    x[i[ii - 1]] = storage[ii] - storage[ii - 1]
  end
  # Compatibility with Convex.jl
  val = zeros(1, 1)
  val[1] = dot(x, y)
  return val
end

function get_cv(F::LovaszExtAbsAtom)
  variable = []
  return push!(variable, F.variable)
end

function get_sv(F::LovaszExtAbsAtom)
  return []
end

function get_v(F::LovaszExtAbsAtom)
  variable = []
  return push!(variable, F.variable)
end


module DONEs

using Distributions
using LinearAlgebra
using NLopt

include("RFEs.jl")

export RFE, DONE, add_measurement!, update_optimal_input!, new_input, evaluateRFE

"""
    DONE(
        rfe::RFE # Random Feature/Fourier Expansion

        # Variables
        current_optimal_x::Vector{T} where T <: AbstractFloat
        n::Int # number of variables
        lower_bound::Vector{T} where T <: AbstractFloat
        upper_bound::Vector{T} where T <: AbstractFloat

        # DONE algorithm sliding window variant
        sliding_window::Bool
        sliding_window_length::Int
        past_inputs::Vector{Vector{T}} where T <: AbstractFloat
        past_outputs::Vector{T} where T <: AbstractFloat

        # Algorithm variables
        iteration::Int

        # Exploration
        surrogate_exploration_prob_dist::Distributions.Distribution
        function_exploration_prob_dist::Distributions.Distribution)

A DONE (Data-based Online Non-linear Extremun-seeker) struct is the main structure for finding the minimum of an unknown function that is typically expensive or difficult to evaluate and that returns a
'measurement' that is corrupted by noise.
The purpose of the underlying algorithm is to find the minimum of this function in as few measurements as possible.

For details of the underlying algorithm, see:
L. Bliek, H.R.G.W. Verstraete, M. Verhaegen and S. Wahls - Online otimization with costly and noisy measurements using random Fourier expansions. (https://arxiv.org/abs/1603.09620)
L. Bliek - Automatic Tuning of Photonic Beamformers.

The (unknown) function f(x) that we want to minimize is approximated using a RFE (Random Fourier/Feature Expansion, type `?RFE`).
n = length(x) and lower_bound .< x .< upper_bound

If f changes over time, then a sliding window of measurements can be used.

The struct keeps track of the current estimate of the optimal x. This is computed from the estimated function (based on the RFE) and a minimization using the NLopt package and the L-BFGS algorithm, initialized using the currently estimated optimal x + some perturbation from the surrogate_exploration_prob_dist probability distribution.

Finding the minimum of the function is balanced with staying at the minimum of the function through the use of exploration of the function using the function_exploration_prob_dist probability distribution to suggest a new measurement around the optimal x + some perturbation.
"""
mutable struct DONE
    rfe::RFE # Random Feature/Fourier Expansion

    # Variables
    current_optimal_x::Vector{T} where T <: AbstractFloat
    n::Int # number of variables
    lower_bound::Vector{T} where T <: AbstractFloat
    upper_bound::Vector{T} where T <: AbstractFloat

    # DONE algorithm sliding window variant
    sliding_window::Bool
    sliding_window_length::Int
    past_inputs::Vector{Vector{T}} where T <: AbstractFloat
    past_outputs::Vector{T} where T <: AbstractFloat

    # Algorithm variables
    iteration::Int

    # Exploration
    surrogate_exploration_prob_dist::Distributions.Distribution
    function_exploration_prob_dist::Distributions.Distribution
end

"""
    DONE(rfe, lower_bound, upper_bound, σ_surrogate_exploration, σ_function_exploration; sliding_window=false, sliding_window_length=1 )

The surrogate explaration and the function exploration are initialized with a probability distribution N(0,σI).
sliding_window_length is the number of past measurements to keep into account.
"""
function DONE(rfe, lower_bound, upper_bound, σ_surrogate_exploration, σ_function_exploration; sliding_window=false, sliding_window_length=1 )
    n = size(rfe.Ω,2)
    current_optimal_x = (lower_bound+upper_bound)./2.0
    if sliding_window
        past_inputs = Vector{Float64}[]
        past_outputs = Float64[]
    else
        sliding_window_length=0
        past_inputs = Vector{Float64}[]
        past_outputs = Float64[]
    end
    iteration = 0
    surrogate_exploration_prob_dist = Distributions.MvNormal(zeros(Float64,n),σ_surrogate_exploration*Diagonal(ones(Float64,n)))
    function_exploration_prob_dist = Distributions.MvNormal(zeros(Float64,n),σ_function_exploration*Diagonal(ones(Float64,n)))

    DONE(rfe,current_optimal_x,n,lower_bound,upper_bound,sliding_window,
        sliding_window_length,past_inputs,past_outputs,iteration,
        surrogate_exploration_prob_dist,function_exploration_prob_dist)
end

"""
    add_measurement!(alg::DONE,x::Vector{T} where T <: AbstractFloat,y::AbstractFloat)

Process a new measurement y from a function evaluation f(x).

Implementation details: see page 111 of L. Bliek - Automatic Tuning of Photonic Beamformers
"""
function add_measurement!(alg::DONE,x::Vector{T} where T <: AbstractFloat,y::AbstractFloat)
    @assert length(x) == alg.n

    v = alg.rfe.variable_offset ? alg.rfe.offset : 0.

    # downdate with oldest measurement
    if alg.sliding_window && alg.sliding_window_length + 1 <= alg.iteration
        a, g = downdateRFE!(alg.rfe, alg.past_inputs[1], alg.past_outputs[1]+v)
        if alg.rfe.variable_offset
            alg.rfe.h[:] = alg.rfe.h - g*(1.0 - dot(a,alg.rfe.h))
        end

        # update list of inputs and measurements
        alg.past_inputs[:] = vcat(alg.past_inputs[2:end],[x])
        alg.past_outputs[:] = vcat(alg.past_outputs[2:end],y)
    elseif alg.sliding_window && alg.sliding_window_length + 1 > alg.iteration
        alg.past_inputs[:] = push!(alg.past_inputs,x)
        alg.past_outputs[:] = push!(alg.past_outputs,y)
    end

    # account for variable offset if any
    if alg.rfe.variable_offset && y + v > 0
        alg.rfe.offset = -2y
        alg.rfe.c[:] = alg.rfe.c + (-2y - v)*alg.rfe.h
        v = -2y
    end

    # update with newest measurement
    a,g = updateRFE!(alg.rfe,x,y+v)

    # account for variable offset if any
    if alg.rfe.variable_offset
        alg.rfe.h[:] = alg.rfe.h + g*(1.0 - dot(a,alg.rfe.h))
    end

    alg.iteration = alg.iteration + 1
end

function project_on_bounds(x,lb,ub)
    return [min(max(xi,lbi),ubi) for (xi,lbi,ubi) in zip(x,lb,ub)]
end

"""
    x_opt = update_optimal_input!(alg::DONE)

Based on the current estimate of the function to be minimized, compute (estimate of) the minimizing argument.
"""
function update_optimal_input!(alg::DONE)
    Ω = alg.rfe.Ω
    b = alg.rfe.b
    c = alg.rfe.c
    o = alg.rfe.offset

    f(x) = dot(c, cos.(Ω*x + b)) + o
    ∇f(x) = -Ω' * Diagonal(sin.(Ω*x + b)) * c
    function myfunc(x::Vector,grad::Vector)
        grad[:] = ∇f(x)
        return f(x)
    end

    opt = NLopt.Opt(:LD_LBFGS,alg.n)
    opt.lower_bounds = alg.lower_bound
    opt.upper_bounds = alg.upper_bound
    opt.min_objective = myfunc

    x0 = project_on_bounds(alg.current_optimal_x + rand(alg.surrogate_exploration_prob_dist),alg.lower_bound,alg.upper_bound)
    (minf,minx,ret) = NLopt.optimize(opt, x0)
    alg.current_optimal_x[:] = minx
    return minx
end

"""
    x_new = new_input(alg::DONE)

Generate a preffered new measurement to be done on the (noisy) functino f.
"""
function new_input(alg::DONE)
    return project_on_bounds(alg.current_optimal_x + rand(alg.function_exploration_prob_dist),alg.lower_bound,alg.upper_bound)
end

# Added as reference.
function minimize_DONE(f :: Function, N :: Int64, lb :: Vector{Float64}, ub :: Vector{Float64}, σs :: Float64 = 0.1, σf :: Float64 = 0.1)
    n_basis = 3
    rfe = RFE(20, n_basis, 0.1)
    done = DONE(rfe, lb, ub, σs, σf)
    best_x :: Union{Nothing, Vector{Float64}} = nothing
    best_y :: Float64 = typemax(Float64)
    for i in 1:N
        xi = new_input(done)
        yi = f(xi)
        if yi < best_y
            best_x = xi
            best_y = yi
        end
        add_measurement!(done, xi, yi)
        update_optimal_input!(done)
    end
    return best_y, best_x
end

end # module


include(joinpath("..", "common", "utils.jl"))
include("cg_sem.jl")
import GaussQuadrature
using PGFPlotsX: @pgf, Plot, Axis, pgfsave, Table, LegendEntry, GroupPlot

Base.@kwdef struct Params{T}
    T_virt_surf::T = 280
    T_min_ref::T = 230
    MSLP::T = 10^5
    grav::T = 98 // 100
    R_d::T = 287
    γ::T = 7 // 5
end

# Taken from Daniel's code
function decaying_temperature_profile(params::Params, z)
    # Scale height for surface temperature
    H_sfc = params.R_d * params.T_virt_surf / params.grav
    H_t = H_sfc

    z′ = z / H_t
    tanh_z′ = tanh(z′)

    ΔTv = params.T_virt_surf - params.T_min_ref
    Tv = params.T_virt_surf - ΔTv * tanh_z′

    ΔTv′ = ΔTv / params.T_virt_surf
    p = -H_t * (z′ + ΔTv′ * (log(1 - ΔTv′ * tanh_z′) - log(1 + tanh_z′) + z′))
    p /= H_sfc * (1 - ΔTv′^2)
    p = params.MSLP * exp(p)
    ρ = p / (params.R_d * Tv)
    return (p = p, ρ = ρ)
end

function plot_comparison(output_file, ρ, ρ_ref, ρe, ρe_ref, z)
    # Plot ρ_ref and ρ
    @pgf ρ_plot = Axis()
    @pgf push!(ρ_plot, Plot({color = "red", no_marks}, Table(ρ_ref[:], z[:])))
    @pgf push!(ρ_plot, LegendEntry(raw"$\rho_{ref}$"))
    @pgf push!(ρ_plot, Plot({color = "blue", no_marks}, Table(ρ[:], z[:])))
    @pgf push!(ρ_plot, LegendEntry(raw"$\rho_{hb}$"))

    # Plot ρ_ref - ρ
    @pgf Δρ_plot = Axis()
    @pgf push!(
        Δρ_plot,
        Plot({color = "red", no_marks}, Table(ρ_ref[:] - ρ[:], z[:])),
    )
    @pgf push!(Δρ_plot, LegendEntry(raw"$\rho_{ref} - \rho_{hb}$"))

    # Plot ρe_ref and ρe
    @pgf ρe_plot = Axis()
    @pgf push!(ρe_plot, Plot({color = "red", no_marks}, Table(ρe_ref[:], z[:])))
    @pgf push!(ρe_plot, LegendEntry(raw"$\rho e_{ref}$"))
    @pgf push!(ρe_plot, Plot({color = "blue", no_marks}, Table(ρe[:], z[:])))
    @pgf push!(ρe_plot, LegendEntry(raw"$\rho e_{hb}$"))

    # Plot ρe_ref - ρe
    @pgf ΔE_plot = Axis()
    @pgf push!(
        ΔE_plot,
        Plot(
            {color = "red", no_marks},
            Table((ρe_ref[:] - ρe[:]) ./ (ρe_ref[:]), z[:]),
        ),
    )
    @pgf push!(
        ΔE_plot,
        LegendEntry(raw"$(\rho e_{ref} - \rho e_{hb})/\rho e_{ref}$"),
    )

    # Make subplots
    @pgf hb_compare = GroupPlot(
        {
            group_style = {
                group_size = "2 by 2",
                horizontal_sep = "50pt",
                vertical_sep = "50pt",
            },
            legend_pos = "north east",
            ymin = minimum(z[:]),
            ymax = maximum(z[:]),
            ylabel = "height",
            # "scaled y ticks = false",
            # "scaled x ticks = false",
        },
        ρ_plot,
        ρe_plot,
        Δρ_plot,
        ΔE_plot,
    )

    pgfsave(output_file, hb_compare)
end

#=
let
    T = Float64
    N = 5
    Nq = N + 1
    K = 10
    #
    z0, z1 = 0, 30e3
    elem = element_operators(N)
    mesh = create_mesh(K, elem, z0, z1)
    #
    p = rand(T, Nq, K)
    #
    dg2cg_scatter(p, mesh, elem)
    #
    params = Params{T}()
    #
    ρ0 = -elem.D * p / (mesh.J * params.grav)
    #
    ρ = balanced_rho(p, params.grav, mesh, elem)
    #
    Pr = DG_cutoff_filter_matrix(elem.ξ, N)
    @assert all(Pr * ρ .≈ ρ0)
    ρ1 = dg2cg_scatter(copy(ρ), mesh, elem)
    @assert all(ρ .≈ ρ1)
end

let
    T = Float64
    N = 4
    Nq = N + 1
    K = 250
    #
    z0, z1 = 0, 30e4
    elem = element_operators(N)
    mesh = create_mesh(K, elem, z0, z1)
    #
    params = Params{T}()
    #
    p, ρref = similar(mesh.zdg), similar(mesh.zdg)
    for (i, z) in enumerate(mesh.zdg)
        (p[i], ρref[i]) = decaying_temperature_profile(params, z)
    end
    #
    ρ_dg = -elem.D * p / (mesh.J * params.grav)
    #
    ρ_cg = balanced_rho(p, params.grav, mesh, elem)
    Pr = DG_cutoff_filter_matrix(elem.ξ, N)
    @assert all(Pr * ρ_cg .≈ ρ_dg)
    #
    nothing
end
=#

let
    T = Float64
    N = 4
    Nq = N + 1
    K = 15
    #
    z0, z1 = 0, 30e4
    elem = element_operators(N)
    mesh = create_mesh(K, elem, z0, z1)
    #
    params = Params{T}()
    #
    p, ρ_ref = similar(mesh.zdg), similar(mesh.zdg)
    for (i, z) in enumerate(mesh.zdg)
        (p[i], ρ_ref[i]) = decaying_temperature_profile(params, z)
    end
    #
    ρ = balanced_rho(p, params.grav, mesh, elem)
    #
    w = fill!(similar(ρ), 0)
    #
    ρe_ref =
        ρ_ref .* w .^ 2 / 2 +
        p / (params.γ - 1) +
        params.grav * ρ_ref .* mesh.zdg
    ρe = ρ .* w .^ 2 / 2 + p / (params.γ - 1) + params.grav * ρ .* mesh.zdg
    #
    #=
    plot_comparison(
        "balance_compare_initial.pdf",
        ρ,
        ρ_ref,
        ρe,
        ρe_ref,
        mesh.zdg,
    )
    #
    ∂q = (
        ρ = fill!(similar(ρ), 0),
        w = fill!(similar(w), 0),
        ρe = fill!(similar(ρe), 0),
    )
    =#
    pres(ρ, w, ρe, z) =
        (params.γ - 1) * (ρe - ρ * w^2 / 2 - ρ * params.grav * z)
    #
    #
    # q = (ρ = ρ_ref, w = w, ρe = ρe_ref)
    q = (ρ = ρ, w = w, ρe = ρe)
    q0 = (ρ = copy(q.ρ), w = copy(q.w), ρe = copy(q.ρe))
    Pr = DG_cutoff_filter_matrix(elem.ξ, N)
    #=
    q_ref = (ρ = ρ_ref, w = w, ρe = ρe_ref)
    #
    element_tendency!(∂q, q_ref, elem, mesh.J, mesh.zdg, params.grav, pres)
    println("Using reference values without projection")
    @show extrema(∂q.w)
    println()
    #
    Pr = DG_cutoff_filter_matrix(elem.ξ, N)
    element_tendency!(∂q, q_ref, elem, mesh.J, mesh.zdg, params.grav, pres, Pr)
    println("Using reference values with projection")
    @show extrema(∂q.w)
    println()
    #
    #
    element_tendency!(∂q, q, elem, mesh.J, mesh.zdg, params.grav, pres)
    println("Using correct values without projection")
    @show extrema(∂q.w)
    println()
    #
    element_tendency!(∂q, q, elem, mesh.J, mesh.zdg, params.grav, pres, Pr)
    println("Using correct values with projection")
    @show extrema(∂q.w)
    println()
    =#
    #
    c = maximum(wave_speed.(ρ, pres.(q.ρ, q.w, q.ρe, mesh.zdg), params.γ))
    #
    dt = mesh.Δz / (N^2 * c)
    #
    @show extrema(q.ρ)
    rhs!(∂q, q, t) = tendency!(∂q, q, t, mesh, elem, params.grav, pres, Pr)
    ndays = 1
    tspan = (0, 86400 * ndays)
    tspan = (0, 100 * dt)
    timestep!(q, rhs!, dt, tspan)
    #
    plot_comparison("balanced.pdf", q0.ρ, q.ρ, q0.ρe, q.ρe, mesh.zdg)
    #
    @show extrema(q0.ρ)
    @show extrema(q.ρ)
    @show extrema(q.ρ - q0.ρ)
    @show extrema(q.w)
    @show sum(elem.ω .* q0.ρ) - sum(elem.ω .* q.ρ)
end
nothing