# # Perfect convective adjustment calibration with Unscented Ensemble Kalman Inversion
#
# This example calibrates a convective adjustment model in the "perfect model context".
# In this context, synthetic observations are generated by a convective adjustment model
# with "true" parameters. The true parameters are then "rediscovered" by calibrating the model
# to match the synthetic observations.
#
# We use the discrepency between observed and modeled buoyancy ``b`` to calibrate
# the convective adjustment model.
# The calibration problem is solved by Ensemble Kalman Inversion.
# For more information about Ensemble Kalman Inversion, see the
# [EnsembleKalmanProcesses.jl documentation](https://clima.github.io/EnsembleKalmanProcesses.jl/stable/ensemble_kalman_inversion/).
#
# The calibration is done here using Unscented Kalman Inversion. For more information about the 
# algorithm refer to [EnsembleKalmanProcesses.jl documentation](https://clima.github.io/EnsembleKalmanProcesses.jl/stable/unscented_kalman_inversion/).

# ## Install dependencies
#
# First let's make sure we have all required packages installed.

# ```julia
# using Pkg
# pkg"add OceanTurbulenceParameterEstimation, Oceananigans, Distributions, CairoMakie"
# ```

# We load some packages,

using OceanTurbulenceParameterEstimation
using LinearAlgebra
using EnsembleKalmanProcesses.ParameterDistributionStorage

# and reuse some some code from a previous example to generate observations,

examples_path = joinpath(pathof(OceanTurbulenceParameterEstimation), "..", "..", "examples")
include(joinpath(examples_path, "intro_to_inverse_problems.jl"))

data_path = generate_synthetic_observations()
observations = SyntheticObservations(data_path, field_names=:b, normalize=ZScore)


# ### Free parameters
#
# We construct some prior distributions for our free parameters. We found that it often helps to
# constrain the prior distributions so that neither very high nor very low values for diffusivities
# can be drawn out of the distribution.

priors = (
    convective_κz = ConstrainedNormal(0.0, 1.0, 0.0, 4 * θ★.convective_κz),
    background_κz = ConstrainedNormal(0.0, 1.0, 0.0, 4 * θ★.background_κz)
)

free_parameters = FreeParameters(priors)

# and an ensemble_simulation,

Nparameters = length(priors)
Nensemble = 2 * Nparameters + 1

ensemble_simulation, closure★ = build_ensemble_simulation(observations; Nensemble)

# To visualize the prior distributions we randomly sample out values from then and plot the p.d.f.

using CairoMakie
using OceanTurbulenceParameterEstimation.EnsembleKalmanInversions: convert_prior, inverse_parameter_transform

samples(prior) = [inverse_parameter_transform(prior, x) for x in rand(convert_prior(prior), 50000000)]

samples_convective_κz = samples(priors.convective_κz)
samples_background_κz = samples(priors.background_κz)

f = Figure()

axtop = Axis(f[1, 1],
             xlabel = "convective_κz [m² s⁻¹]",
             ylabel = "p.d.f.")

axbottom = Axis(f[2, 1],
                xlabel = "background_κz [m² s⁻¹]",
                ylabel = "p.d.f.")

densities = []

push!(densities, density!(axtop, samples_convective_κz))
push!(densities, density!(axbottom, samples_background_κz))

save("visualize_prior_diffusivities_convective_adjustment_uki.svg", f); nothing #hide 

# ![](visualize_prior_diffusivities_convective_adjustment_uki.svg)

# ### The inverse problem

# We can construct the inverse problem ``y = G(θ) + η``. Here, ``y`` are the `observations` and ``G`` is the
# `ensemble_model`.
calibration = InverseProblem(observations, ensemble_simulation, free_parameters)

# Next, we construct an `UnscentedKalmanInversion` (UKI) object,

prior_mean = fill(0.0, Nparameters) 
prior_cov = Matrix(Diagonal(fill(1.0, Nparameters)))
α_reg = 1.0   # regularization parameter 
update_freq = 1
noise_covariance = 0.05^2  # error is about 5%

uki = UnscentedKalmanInversion(calibration, prior_mean, prior_cov;
                               noise_covariance = noise_covariance, α_reg = α_reg, update_freq = update_freq)

# and perform few iterations to see if we can converge to the true parameter values.

iterate!(uki; iterations = 10)

# Last, we visualize the outputs of UKI calibration:

θ_mean, θθ_cov, θθ_std_arr, error =  UnscentedKalmanInversionPostprocess(uki)

N_iter = size(θ_mean, 2)

f = Figure(resolution = (800, 600))
ax1 = Axis(f[1, 1],
           xlabel = "iterations",
           xticks = 1:N_iter,
           ylabel = "convective_κz [m² s⁻¹]")

ax2 = Axis(f[2, 1],
           xlabel = "iterations",
           xticks = 1:N_iter,
           ylabel = "background_κz [m² s⁻¹]")

band!(ax1, 1:N_iter, θ_mean[1, :] .+ θθ_std_arr[1, :], θ_mean[1, :] .- θθ_std_arr[1, :])
lines!(ax1, Float64.(1:N_iter), θ_mean[1, :])
hlines!(ax1, [θ★.convective_κz], color=:red)

band!(ax2, 1:N_iter, θ_mean[2, :] .+ θθ_std_arr[2, :], θ_mean[2, :] .- θθ_std_arr[2, :])
lines!(ax2, 1:N_iter, θ_mean[2, :])
hlines!(ax2, [θ★.background_κz], color=:red)

xlims!(ax1, 0.5, N_iter+0.5)
xlims!(ax2, 0.5, N_iter+0.5)

save("uki_results.svg", f); nothing #hide 

# ![](uki_results.svg)

# and the error:

f = Figure(resolution = (800, 300))

lines(f[1, 1], 2:N_iter, error, color = :red, linewidth = 2,
               axis = (xlabel = "iterations",
                       xticks = 1:N_iter,
                       ylabel = "error"))

xlims!(0.5, N_iter+0.5)

save("uki_error.svg", f); nothing #hide 

# ![](uki_error.svg)
