using CSV, DataFrames  # data loading
using EasyGPs  # handles all things related to GPs
using Plots  # visualisation

(xtrain, ytrain), (xtest, ytest) = let
    data = CSV.read(joinpath("/home/runner/work/EasyGPs.jl/EasyGPs.jl/examples/0-mauna-loa", "CO2_data.csv"), Tables.matrix; header=0)
    year = data[:, 1]
    co2 = data[:, 2]

    # We split the data into training and testing set:
    idx_train = year .< 2004
    idx_test = .!idx_train

    (year[idx_train], co2[idx_train]), (year[idx_test], co2[idx_test])  # block's return value
end

function plotdata()
    plot(; xlabel="year", ylabel="CO₂ [ppm]", legend=:bottomright)
    scatter!(xtrain, ytrain; label="training data", ms=2, markerstrokewidth=0)
    return scatter!(xtest, ytest; label="test data", ms=2, markerstrokewidth=0)
end

plotdata()

k_smooth_trend = exp(8.0) * with_lengthscale(SEKernel(), exp(4.0))#with_lengthscale(SEKernel(), exp(4.0))
k_seasonality =
    exp(2.0) * PeriodicKernel(; r=[0.5]) * with_lengthscale(SEKernel(), exp(4.0))
k_medium_term_irregularities =
    1.0 * with_lengthscale(RationalQuadraticKernel(; α=exp(-1.0)), 1.0)
k_noise_terms =
    exp(-4.0) * with_lengthscale(SEKernel(), exp(-2.0)) + exp(-4.0) * WhiteKernel()
kernel = k_smooth_trend + k_seasonality + k_medium_term_irregularities + k_noise_terms

gp = GP(kernel)

fpost_init = posterior(gp(xtrain), ytrain)

plot_gp!(f; label) = plot!(f(1920:0.2:2030); ribbon_scale=2, linewidth=1, label)

plotdata()
plot_gp!(fpost_init; label="posterior f(⋅)")

@time fitted_gp = EasyGPs.fit(
    gp,
    xtrain,
    ytrain;
    optimizer=Optim.LBFGS(;
        alphaguess=Optim.LineSearches.InitialStatic(; scaled=true),
        linesearch=Optim.LineSearches.BackTracking(),
    ),
)

fpost_opt = posterior(fitted_gp(xtrain), ytrain)

fpost_opt.prior.kernel

plotdata()
plot_gp!(fpost_opt; label="optimized posterior f(⋅)")