text:
########################################################################################
# Caches
########################################################################################
abstract type ODEFiltersCache <: OrdinaryDiffEq.OrdinaryDiffEqCache end
mutable struct GaussianODEFilterCache{
    RType,
    ProjType,
    SolProjType,
    PType,
    PIType,
    EType,
    uType,
    duType,
    xType,
    AType,
    QType,
    matType,
    diffusionType,
    diffModelType,
    measType,
    puType,
    llType,
    CType,
    rateType,
    UF,
    JC,
    uNoUnitsType,
} <: ODEFiltersCache
    # Constants
    d::Int                  # Dimension of the problem
    q::Int                  # Order of the prior
    A::AType
    Q::QType
    Ah::AType
    Qh::QType
    diffusionmodel::diffModelType
    R::RType
    Proj::ProjType
    SolProj::SolProjType
    # Also mutable
    P::PType
    PI::PIType
    E0::EType
    E1::EType
    E2::EType
    # Mutable stuff
    u::uType
    u_pred::uType
    u_filt::uType
    tmp::uType
    atmp::uNoUnitsType
    x::xType
    x_pred::xType
    x_filt::xType
    x_tmp::xType
    x_tmp2::xType
    measurement::measType
    m_tmp::measType
    pu_tmp::puType
    H::matType
    du::duType
    ddu::matType
    K1::matType
    K2::matType
    G1::matType
    G2::matType
    covmatcache::matType
    default_diffusion::diffusionType
    local_diffusion::diffusionType
    global_diffusion::diffusionType
    err_tmp::duType
    log_likelihood::llType
    C1::CType
    C2::CType
    du1::rateType
    uf::UF
    jac_config::JC
end

function OrdinaryDiffEq.alg_cache(
    alg::GaussianODEFilter,
    u,
    rate_prototype,
    uEltypeNoUnits,
    uBottomEltypeNoUnits,
    tTypeNoUnits,
    uprev,
    uprev2,
    f,
    t,
    dt,
    reltol,
    p,
    calck,
    ::Val{IIP},
) where {IIP}
    initialize_derivatives = true

    if u isa Number
        error("We currently don't support scalar-valued problems")
    end

    is_secondorder_ode = f isa DynamicalODEFunction
    if is_secondorder_ode
        @warn "Assuming that the given ODE is a SecondOrderODE. If this is not the case, e.g. because it is some other dynamical ODE, the solver will probably run into errors!"
    end

    q = alg.order
    d = is_secondorder_ode ? length(u[1, :]) : length(u)
    D = d * (q + 1)

    u_vec = u[:]
    t0 = t

    uType = typeof(u)
    uElType = eltype(u_vec)
    matType = Matrix{uElType}

    # Projections
    Proj = projection(d, q, uElType)
    E0, E1, E2 = Proj(0), Proj(1), Proj(2)
    @assert f isa AbstractODEFunction
    SolProj = f isa DynamicalODEFunction ? [Proj(1); Proj(0)] : Proj(0)

    # Prior dynamics
    @assert alg.prior == :ibm "Only the ibm prior is implemented so far"
    P, PI = init_preconditioner(d, q, uElType)

    A, Q = ibm(d, q, uElType)

    initial_variance = 1.0 * ones(uElType, D)
    x0 = Gaussian(zeros(uElType, D), SRMatrix(Matrix(Diagonal(sqrt.(initial_variance)))))

    # Measurement model
    R = zeros(uElType, d, d)

    # Pre-allocate a bunch of matrices
    h = zeros(uElType, d)
    H = f isa DynamicalODEFunction ? copy(E2) : copy(E1)
    du = f isa DynamicalODEFunction ? similar(u[2, :]) : similar(u)
    ddu = f isa DynamicalODEFunction ? zeros(uElType, d, 2d) : zeros(uElType, d, d)
    v, S = similar(h), similar(ddu)
    v = similar(h)
    S =
        alg isa EK0 ? SRMatrix(zeros(uElType, d, D), Diagonal(zeros(uElType, d, d))) :
        SRMatrix(zeros(uElType, d, D), zeros(uElType, d, d))
    measurement = Gaussian(v, S)
    pu_tmp =
        f isa DynamicalODEFunction ?
        Gaussian(zeros(uElType, 2d), SRMatrix(zeros(uElType, 2d, D))) : similar(measurement)
    K = zeros(uElType, D, d)
    G = zeros(uElType, D, D)
    C1 = SRMatrix(zeros(uElType, D, 2D), zeros(uElType, D, D))
    C2 = SRMatrix(zeros(uElType, D, 3D), zeros(uElType, D, D))
    covmatcache = similar(G)

    if alg isa EK1FDB
        H = [E1; E2]
        v = [v; v]
        S = SRMatrix(zeros(uElType, 2d, D), zeros(uElType, 2d, 2d))
        measurement = Gaussian(v, S)
        K = zeros(uElType, D, 2d)
    end

    diffmodel = alg.diffusionmodel
    initdiff = initial_diffusion(diffmodel, d, q, uEltypeNoUnits)
    copy!(x0.Σ, apply_diffusion(x0.Σ, initdiff))

    Ah, Qh = copy(A), copy(Q)
    u_pred = similar(u)
    u_filt = similar(u)
    tmp = similar(u)
    x_pred = copy(x0)
    x_filt = similar(x0)
    x_tmp = similar(x0)
    x_tmp2 = similar(x0)
    m_tmp = similar(measurement)
    K2 = similar(K)
    G2 = similar(G)
    err_tmp = similar(du)

    # Things for calc_J
    uf =
        IIP == true ? OrdinaryDiffEq.UJacobianWrapper(f, t, p) :
        OrdinaryDiffEq.UDerivativeWrapper(f, t, p)
    du1 = similar(rate_prototype)
    dw1 = zero(u)
    jac_config = OrdinaryDiffEq.build_jac_config(alg, f, uf, du1, uprev, u, tmp, dw1)
    atmp = similar(u, uEltypeNoUnits)

    return GaussianODEFilterCache{
        typeof(R),
        typeof(Proj),
        typeof(SolProj),
        typeof(P),
        typeof(PI),
        typeof(E0),
        uType,
        typeof(du),
        typeof(x0),
        typeof(A),
        typeof(Q),
        matType,
        typeof(initdiff),
        typeof(diffmodel),
        typeof(measurement),
        typeof(pu_tmp),
        uEltypeNoUnits,
        typeof(C1),
        typeof(du1),
        typeof(uf),
        typeof(jac_config),
        typeof(atmp),
    }(
        # Constants
        d,
        q,
        A,
        Q,
        Ah,
        Qh,
        diffmodel,
        R,
        Proj,
        SolProj,
        P,
        PI,
        E0,
        E1,
        E2,
        # Mutable stuff
        u,
        u_pred,
        u_filt,
        tmp,
        atmp,
        x0,
        x_pred,
        x_filt,
        x_tmp,
        x_tmp2,
        measurement,
        m_tmp,
        pu_tmp,
        H,
        du,
        ddu,
        K,
        K2,
        G,
        G2,
        covmatcache,
        initdiff,
        initdiff * NaN,
        initdiff * NaN,
        err_tmp,
        zero(uEltypeNoUnits),
        C1,
        C2,
        du1,
        uf,
        jac_config,
    )
end


text:
using Pkg
Pkg.activate("examples")
Pkg.instantiate()

# for std()
using Statistics

# for det()
using LinearAlgebra

# for importing the data from CSV
using DataFrames
using CSV

# For GLOM
import GPLinearODEMaker; GLOM = GPLinearODEMaker

# For this module
using GLOM_RV_Example; GLOM_RV = GLOM_RV_Example

# For units in orbit fitting functions
using UnitfulAstro, Unitful

# for PCA and std()
using MultivariateStats
using Statistics

using Plots
using JLD2

#####################################################################
# CTRL+F "CHANGE"TO FIND PLACES WHERE YOU SHOULD MAKE MODIFICATIONS #
#####################################################################

## Problem setup
length(ARGS) > 0 ? n_bs = parse(Int, ARGS[1]) : n_bs = 5
@assert 1 < n_bs < 6

# CHANGE: importing EXPRES data
length(ARGS) > 1 ? star_ind = parse(Int, ARGS[2]) : star_ind = 4
stars = ["101501", "10700", "26965", "34411"]
star_str = stars[star_ind]
data_dir = "examples/"
data = CSV.read(data_dir * star_str * "_Yale_Wisconsin_hgrvsafe_results.csv", DataFrame)

# CHANGE: the stars rotation rate which is used as the first guess for some GLOM
# hyperparameters and starting point for priors
initial_hypers = [17.1, [0.0069, 34, 1], 37., [0.0069, 22, 1]][star_ind]  # days

# CHANGE: choose a kernel, I suggest 3 for Matern 5/2 or 4 for Quasi-periodic
# kernel
if length(initial_hypers) == 1
    # kernel_names = ["pp", "se", "m52", "qp", "m52_m52", "se_se"]
    kernel_name = "m52"
elseif length(initial_hypers) == 3
    kernel_name = "m52_m52"
end
kernel_function, num_kernel_hyperparameters = GLOM.include_kernel(kernel_name)

# CHANGE: observation times go here
if star_str == "101501"
    obs_xs = collect(data[!, "TELMJD"])
else
    obs_xs = collect(data[!, "BARYMJD"])
end
sort_inds = sortperm(obs_xs)
obs_xs = obs_xs[sort_inds]

# taking out the mean observation times makes the optimization easier for models
# with periodic parameters
GLOM_RV.remove_mean!(obs_xs)

# CHANGE: Where to save intermediate data products
save_dir = data_dir * "data/" * star_str * "_"
fig_dir = data_dir * "figs/" * star_str * "_"

# CHANGE: rvs and their errors go here
obs_rvs = data[!, "RV"][sort_inds]
# inject_ks = GLOM_RV.kep_signal(; K=50u"m/s", P=sqrt(2)*5u"d", M0=rand()*2*π, ω_or_k=rand()*2*π, e_or_h=0.1)
# obs_rvs[:] .+= ustrip.(inject_ks.(obs_xs.*u"d"))
GLOM_RV.remove_mean!(obs_rvs)
obs_rvs_err = data[!, "RV_std"][sort_inds]

# CHANGE: whether to use saved values if they exist
use_saved = true

## PCA for indicators

use_PCA = false
use_PPCA = false

find_data_col_ind(col_name) = findfirst(x -> x==col_name, names(data))

if use_PCA || use_PPCA
    b_inds = find_data_col_ind("b0"):(find_data_col_ind("b0") + n_bs - 1)
    t_inds = find_data_col_ind("t0") - find_data_col_ind("b0") .+ b_inds
    Xte = Matrix(data[!, b_inds])'[:, sort_inds]

    if use_PPCA
        #PPCA version
        M = fit(PPCA, Xte; maxoutdim=2)
        # modifed from transform() at from https://github.com/JuliaStats/MultivariateStats.jl/blob/master/src/ppca.jl
        function transforming_mat(m::PPCA{T}) where {T<:Real}
            W  = m.W
            n = outdim(m)
            M = W'W .+ m.σ² * Matrix{T}(I, n, n)
            return inv(M)*m.W'
        end
        Z = transforming_mat(M)
    elseif use_PCA
        # PCA version
        M = fit(PCA, Xte; maxoutdim=2)
        Z = M.proj'
    end
    println("Z: ", Z)
    z = MultivariateStats.transform(M, Xte)  # = Z * (Xte .- M.mean)
    var_β = ((Matrix(data[!, b_inds])[sort_inds, :] ./ Matrix(data[!, t_inds][sort_inds, :]))').^2
    σ_z = zeros(size(z))
    for i in 1:size(z, 2)
        σ_z[:, i] = sqrt.(diag(Z * Diagonal(var_β[:, 1]) * Z'))
    end
else
    b_inds = find_data_col_ind("b2"):find_data_col_ind("b3")
    t_inds = find_data_col_ind("t2"):find_data_col_ind("t3")
    z = Matrix(data[!, b_inds])'[:, sort_inds]
    σ_z = z ./ Matrix(data[!, t_inds])'[:, sort_inds]
end

## Adding indicators

# initializing RV and indicator array
rvs_and_inds_holder = [obs_rvs]
rvs_and_inds_err_holder = [obs_rvs_err]

function append_vec!(to_be_append_to::Vector{Vector{T}}, to_append::Vector{T}; remove_mean::Bool=true) where T
    if remove_mean
        GLOM_RV.remove_mean!(to_append)
    end
    append!(to_be_append_to, [to_append])
end
function add_indicator!(rvs_and_inds::Vector{Vector{T}}, rvs_and_inds_err::Vector{Vector{T}}, indicator::Vector{T}, indicator_err::Vector{T}) where T
    append_vec!(rvs_and_inds, indicator; remove_mean=true)
    append_vec!(rvs_and_inds_err, indicator_err; remove_mean=false)
end

# CHANGE: activity indicators and their errors go here
# you can actually have as many as you want, but obviously it will take longer
# to fit
for i in 1:size(z, 1)
    add_indicator!(rvs_and_inds_holder, rvs_and_inds_err_holder, z[i, :], σ_z[i, :])
end

# CHANGE: which of the indicators do you want use?
inds_to_use = 1:length(rvs_and_inds_holder)
@assert inds_to_use[1] == 1  # inds_to_use needs to include the RVs first

rvs_and_inds = rvs_and_inds_holder[inds_to_use]
rvs_and_inds_err = rvs_and_inds_err_holder[inds_to_use]
n_out = length(rvs_and_inds)  # number of indicators + 1
obs_ys = collect(Iterators.flatten(zip(rvs_and_inds...)))
obs_noise = collect(Iterators.flatten(zip(rvs_and_inds_err...)))

# How many differention orders we want in the GLOM model
n_dif = 2 + 1

# CHANGE: consider changing a (the GLOM coefficients that are used, see
# commented lines below)
# If all a's active:
glo = GLOM.GLO(kernel_function, num_kernel_hyperparameters, n_dif, n_out, obs_xs, copy(obs_ys); noise=copy(obs_noise), a=ones(n_out, n_dif))
# glo = GLOM.GLO(kernel_function, num_kernel_hyperparameters, n_dif, n_out, obs_xs, copy(obs_ys); noise=copy(obs_noise), a=[[1. 1 0];[1 0 1];[1 0 1]])

# Makes the std of each output equal to 1, improves fitting stability
# the normalizations are stored in glo.normals
GLOM.normalize_GLO!(glo)

# CHANGE: Setting initial fit values
initial_total_hyperparameters = collect(Iterators.flatten(glo.a))
append!(initial_total_hyperparameters, initial_hypers)

## Fitting GLOM Model

# CHANGE: Setting kernel hyperparameter priors and kick function
# kick functions help avoid saddle points
tighten_lengthscale_priors = 6
if kernel_name in ["pp", "se", "m52"]
    kernel_hyper_priors(hps::Vector{<:Real}, d::Integer) =
        GLOM_RV.kernel_hyper_priors_1λ(hps, d, initial_hypers, initial_hypers / tighten_lengthscale_priors)
    add_kick!(hps::Vector{<:Real}) = GLOM_RV.add_kick_1λ!(hps)
elseif kernel_name == "qp"
    kernel_hyper_priors(hps::Vector{<:Real}, d::Integer) =
        GLOM_RV.kernel_hyper_priors_qp(hps, d, initial_hypers, initial_hypers ./ tighten_lengthscale_priors)
    add_kick!(hps::Vector{<:Real}) = GLOM_RV.add_kick_qp!(hps)
elseif kernel_name in ["se_se", "m52_m52"]
    kernel_hyper_priors(hps::Vector{<:Real}, d::Integer) =
        GLOM_RV.kernel_hyper_priors_2λ(hps, d, initial_hypers, initial_hypers ./ tighten_lengthscale_priors)
    add_kick!(hps::Vector{<:Real}) = GLOM_RV.add_kick_2λ!(hps)
else
    # kernel_hyper_priors(hps::Vector{<:Real}, d::Integer) = custom function
end

if use_saved && isfile(save_dir*"fit1.jld2")
    # GLOM.include_kernel(kernel_name)
    @load save_dir*"fit1.jld2" glo fit1_total_hyperparameters
else
    # fit_GLOM returns a vector of num_kernel_hyperparameters gp hyperparameters
    # followed by the GLOM coefficients and the Optim result object
    fit1_total_hyperparameters, result = GLOM_RV.fit_GLOM(glo, initial_total_hyperparameters, kernel_hyper_priors, add_kick!)
    println(result)

    @save save_dir*"fit1.jld2" glo fit1_total_hyperparameters
end

## Plotting initial results
# plot_xs = sort(append!(collect(LinRange(-169.73, -169.69, 300)), obs_xs))
# GLOM_at_plot_xs, GLOM_err_at_plot_xs, GLOM_at_obs_xs = GLOM_RV.GLOM_posteriors(glo, plot_xs, fit1_total_hyperparameters)
# scatter(obs_xs, obs_rvs; xrange=(-169.73, -169.69))
# plot!(plot_xs, GLOM_at_plot_xs[1]; ribbons=GLOM_err_at_plot_xs)

plot_xs = sort(append!(collect((obs_xs[1] - 10):3:(obs_xs[end] + 10)), append!([(obs_xs[i] + obs_xs[i - 1]) / 2 for i in 2:length(obs_xs)], obs_xs)))

include("plot_helpers.jl")
plot_helper(prefix::String, fit_total_hyperparameters::AbstractVector) =
    plot_helper(glo, plot_xs, rvs_and_inds, rvs_and_inds_err, prefix, fit_total_hyperparameters)
plot_helper(prefix::String, ks::GLOM_RV.KeplerSignal, fit_total_hyperparameters::AbstractVector) =
    plot_helper(glo_rv, plot_xs, rvs_and_inds, rvs_and_inds_err, prefix, ks, fit_total_hyperparameters)

plot_helper("fit1_", fit1_total_hyperparameters)

## Finding a planet?

nlogprior_hyperparameters(total_hyper::Vector, d::Int) = GLOM_RV.nlogprior_hyperparameters(kernel_hyper_priors, glo.n_kern_hyper, total_hyper, d)
glo_rv = GLO_RV(glo, 1u"d", glo.normals[1]u"m/s")

############
# Post fit #
############

println("starting hyperparameters")
println(initial_total_hyperparameters)
initial_nlogL = GLOM.nlogL_GLOM(glo, initial_total_hyperparameters)
initial_uE = -initial_nlogL - nlogprior_hyperparameters(initial_total_hyperparameters, 0)
println(initial_uE, "\n")

println("ending hyperparameters")
println(fit1_total_hyperparameters)
fit_nlogL1 = GLOM.nlogL_GLOM(glo, fit1_total_hyperparameters)
uE1 = -fit_nlogL1 - nlogprior_hyperparameters(fit1_total_hyperparameters, 0)
println(uE1, "\n")

#########################
# Keplerian periodogram #
#########################

# sample linearly in frequency space so that we get periods from the 1 / uneven Nyquist
freq_grid = GLOM_RV.autofrequency(glo_rv.time; samples_per_peak=11)
period_grid = 1 ./ reverse(freq_grid)
amount_of_periods = length(period_grid)

Σ_obs = GLOM.Σ_observations(glo, fit1_total_hyperparameters)

# making necessary variables local to all workers
fit1_total_hyperparameters_nz = GLOM.remove_zeros(fit1_total_hyperparameters)
nlogprior_kernel = nlogprior_hyperparameters(fit1_total_hyperparameters_nz, 0)

using Distributed

# CHANGE: can change full_fit to false to just do the epicyclic fit which is
# ~40x faster. Distributing the compute only gives me a factor of 2, so only
# try if the keplerian fitting takes a long time
use_distributed = false
full_fit = true

# concurrency is weird so you may have to run this twice
if use_distributed
    GLOM.auto_addprocs()
    @everywhere using Pkg; @everywhere Pkg.activate("examples"); @everywhere Pkg.instantiate()
    @everywhere using GLOM_RV_Example; @everywhere GLOM_RV = GLOM_RV_Example
    @everywhere import GPLinearODEMaker; @everywhere GLOM = GPLinearODEMaker
    @everywhere using UnitfulAstro, Unitful

    GLOM.sendto(workers(), kernel_name=kernel_name)
    @everywhere GLOM.include_kernel(kernel_name)

    GLOM.sendto(workers(), glo_rv=glo_rv, fit1_total_hyperparameters=fit1_total_hyperparameters, Σ_obs=Σ_obs, nlogprior_kernel=nlogprior_kernel, full_fit=full_fit)
end

@everywhere function fit_kep_hold_P(P::Unitful.Time; fast::Bool=false, kwargs...)
    #initialize with fast epicyclic fit
    ks = GLOM_RV.fit_kepler(glo_rv, Σ_obs, GLOM_RV.kep_signal_epicyclic(P=P))
    if !fast
        ks = GLOM_RV.fit_kepler(glo_rv, Σ_obs, GLOM_RV.kep_signal_wright(0u"m/s", P, ks.M0, minimum([ks.e, 0.3]), 0, 0u"m/s"); hold_P=true, avoid_saddle=false, print_stuff=false, kwargs...)
        return ks
    end
    if ks == nothing
        ks = GLOM_RV.fit_kepler(glo_rv, Σ_obs, GLOM_RV.kep_signal_wright(0u"m/s", P, 2 * π * rand(), 0.1, 0, 0u"m/s"); hold_P=true, avoid_saddle=false, print_stuff=false, kwargs...)
        return ks
    end
    return ks
end
@everywhere function kep_unnormalized_posterior_distributed(P::Unitful.Time; kwargs...)
    ks = fit_kep_hold_P(P; kwargs...)
    if ks == nothing
        return [-Inf, -Inf]
    else
        val = GLOM.nlogL_GLOM(
            glo_rv.GLO,
            fit1_total_hyperparameters;
            Σ_obs=Σ_obs,
            y_obs=GLOM_RV.remove_kepler(glo_rv, ks))
        return [-val, GLOM_RV.logprior_kepler(ks; use_hk=false) - nlogprior_kernel - val]
    end
end
@everywhere kep_unnormalized_posterior_distributed(P::Unitful.Time) = kep_unnormalized_posterior_distributed(P; fast=!full_fit)

if use_saved && isfile(save_dir*"period.jld2")
    # need to using Unitful before this will work
    @load save_dir*"period.jld2" likelihoods unnorm_posteriors period_grid best_period
else
    likelihoods = zeros(amount_of_periods)
    unnorm_posteriors = zeros(amount_of_periods)

    @time if use_distributed
        # takes around a minute for 26965 data and 3000 periods
        holder = pmap(x->kep_unnormalized_posterior_distributed(x), period_grid, batch_size=Int(floor(amount_of_periods / (nworkers() + 1)) + 1))
        likelihoods[:] = [holder[i][1] for i in 1:length(holder)]
        unnorm_posteriors[:] = [holder[i][2] for i in 1:length(holder)]
    else
        # takes around minutes for 101501 data and 3000 periods
        for i in 1:amount_of_periods
            likelihoods[i], unnorm_posteriors[i] = kep_unnormalized_posterior_distributed(period_grid[i])
        end
    end

    best_periods = period_grid[GLOM_RV.find_modes(unnorm_posteriors; amount=10)]
    best_period = best_periods[1]

    println("found period:    $(ustrip(best_period)) days")

    @save save_dir*"period.jld2" likelihoods unnorm_posteriors period_grid best_period best_periods
end

plot(ustrip.(period_grid), likelihoods; xaxis=:log, leg=false)
png(fig_dir * "period_lik")
plot(ustrip.(period_grid), unnorm_posteriors; xaxis=:log, leg=false)
png(fig_dir * "period_evi")

####################################################################################################
# Refitting GP with full planet signal at found period subtracted (K,ω,γ-linear, P,M0,e-nonlinear) #
####################################################################################################
remainder(vec, x) = [i > 0 ? i % x : (i % x) + x for i in vec]

# Need to redefine Σ_obs for fit_kep_hold_P()
Σ_obs = GLOM.Σ_observations(glo, fit1_total_hyperparameters)
current_ks = fit_kep_hold_P(best_period; print_stuff=true)
println("before GLOM+Wright fit: ", GLOM_RV.kep_parms_str(current_ks))

if use_saved && isfile(save_dir*"fit3.jld2")
    @load save_dir*"fit3.jld2" glo fit3_total_hyperparameters current_ks
    glo_rv = GLO_RV(glo, 1u"d", glo.normals[1]u"m/s")
else
    # 400s
    @time fit3_total_hyperparameters, current_ks = GLOM_RV.fit_GLOM_and_kep(glo_rv,
    fit1_total_hyperparameters, kernel_hyper_priors, add_kick!, current_ks;
    avoid_saddle=false, fit_alpha=1e-3)

    @save save_dir*"fit3.jld2" glo fit3_total_hyperparameters current_ks
end

# plot_helper("fit3_", current_ks, fit3_total_hyperparameters)

# # these should be near 0
# GLOM_RV.test_∇nlogL_kep(glo_rv, workspace.Σ_obs, current_ks; include_priors=true)
###########################################################################################
# Refitting GP with full planet signal at found period subtracted (K,P,M0,e,ω,γ-nonlinear)#
###########################################################################################

if use_saved && isfile(save_dir*"fit4.jld2")
    @load save_dir*"fit4.jld2" glo fit4_total_hyperparameters current_ks
    glo_rv = GLO_RV(glo, 1u"d", glo.normals[1]u"m/s")
else
    # 200s
    current_ks = GLOM_RV.kep_signal(current_ks)
    println("\nbefore full fit: ", GLOM_RV.kep_parms_str(current_ks))
    @time fit4_total_hyperparameters, current_ks = GLOM_RV.fit_GLOM_and_kep(glo_rv,
    fit3_total_hyperparameters, kernel_hyper_priors, add_kick!, current_ks;
    avoid_saddle=true)

    @save save_dir*"fit4.jld2" glo fit4_total_hyperparameters current_ks
end

plot_helper("fit4_", current_ks, fit4_total_hyperparameters)

full_ks = GLOM_RV.kep_signal(current_ks)

# # these should be near 0
# GLOM_RV.test_∇nlogL_kep(glo_rv, workspace.Σ_obs, full_ks; include_priors=true)

###################
# Post planet fit #
###################

println("fit hyperparameters")
println(fit1_total_hyperparameters)
println(uE1, "\n")

println("kepler hyperparameters")
println(fit4_total_hyperparameters)
fit_nlogL2 = GLOM.nlogL_GLOM(glo, fit4_total_hyperparameters; y_obs=GLOM_RV.remove_kepler(glo_rv, full_ks))
uE2 = -fit_nlogL2 - nlogprior_hyperparameters(fit4_total_hyperparameters, 0) + GLOM_RV.logprior_kepler(full_ks; use_hk=true)
println(uE2, "\n")

println("best fit keplerian")
println(GLOM_RV.kep_parms_str(full_ks))

##################################################################################
# refitting noise model to see if a better model was found during planet fitting #
##################################################################################
if !(use_saved && isfile(save_dir*"fit1.jld2"))
    fit1_total_hyperparameters_temp, result = GLOM_RV.fit_GLOM(
        glo,
        fit4_total_hyperparameters,
        kernel_hyper_priors,
        add_kick!)

    println(result)

    println("first fit hyperparameters")
    println(fit1_total_hyperparameters)
    println(uE1, "\n")

    println("fit after planet hyperparameters")
    println(fit1_total_hyperparameters_temp)
    fit_nlogL1_temp = GLOM.nlogL_GLOM(glo, fit1_total_hyperparameters_temp)
    uE1_temp = -fit_nlogL1_temp - nlogprior_hyperparameters(fit1_total_hyperparameters_temp, 0)
    println(uE1_temp, "\n")

    if uE1_temp > uE1
        println("new fit is better, switching hps")
        fit1_total_hyperparameters[:] = fit1_total_hyperparameters_temp
        fit_nlogL1 = fit_nlogL1_temp
        uE1 = uE1_temp
        @save save_dir*"fit1.jld2" glo fit1_total_hyperparameters
    end
end
##########################
# Evidence approximation #
##########################

# no planet
H1 = (GLOM.∇∇nlogL_GLOM(glo, fit1_total_hyperparameters)
    + nlogprior_hyperparameters(GLOM.remove_zeros(fit1_total_hyperparameters), 2))
try
    global E1 = GLOM.log_laplace_approximation(H1, -uE1, 0)
catch err
    if isa(err, DomainError)
        println("Laplace approximation failed for initial GP fit")
        println("det(H1): $(det(H1)) (should've been positive)")
        global E1 = 0
    else
        rethrow()
    end
end

# planet
H2 = Matrix(GLOM_RV.∇∇nlogL_GLOM_and_planet(glo_rv, fit4_total_hyperparameters, full_ks; include_kepler_priors=true))
n_hyper = length(GLOM.remove_zeros(fit4_total_hyperparameters))
H2[1:n_hyper, 1:n_hyper] += nlogprior_hyperparameters(GLOM.remove_zeros(fit4_total_hyperparameters), 2)
try
    global E2 = GLOM.log_laplace_approximation(Symmetric(H2), -uE2, 0)
catch err
    if isa(err, DomainError)
        println("Laplace approximation failed for planet fit")
        println("det(H2): $(det(H2)) (should've been positive)")
        global E2 = 0
    else
        rethrow()
    end
end

println("\nlog likelihood for GLOM model: " * string(-fit_nlogL1))
println("log likelihood for GLOM + planet model: " * string(-fit_nlogL2))

println("\nunnormalized posterior for GLOM model: " * string(uE1))
println("unnormalized posterior for GLOM + planet model: " * string(uE2))

println("\nevidence for GLOM model: " * string(E1))
println("evidence for GLOM + planet model: " * string(E2))

# # This should be pretty close to true
# GLOM_RV.test_∇∇nlogL_kep(glo_rv, workspace.Σ_obs, current_ks; include_priors=true)

ext: jl


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# Adapted from KrigingEstimators.jl
# ------------------------------------------------------------------

function solve(problem::EstimationProblem, solver::LocalKriging)
  # preprocess user input
  preproc = local_preprocess(problem, solver)

  # results for each variable
  μs = []; σs = []
  for var in name.(variables(problem))
    varμ, varσ = local_solve_approx(problem, var, preproc)
    push!(μs, var => varμ)
    push!(σs, Symbol(var,"-variance") => varσ)
  end

  georef((; μs..., σs...), domain(problem))
end

function local_preprocess(problem::EstimationProblem, solver::LocalKriging)
  # retrieve problem info
  pdomain = domain(problem)
  pdata   = data(problem)
  ndata   = nelements(pdata)

  # result of preprocessing
  preproc = Dict{Symbol,NamedTuple}()

  for covars in covariables(problem, solver)
    for var in covars.names
      # get user parameters
      varparams = covars.params[(var,)]

      # determine which Kriging variant to use
      if varparams.mean ≠ nothing
        estimator = SimpleKriging(varparams.variogram, varparams.mean)
      else
        estimator = OrdinaryKriging(varparams.variogram)
      end

      # determine minimum/maximum number of neighbors
      minneighbors = varparams.minneighbors
      maxneighbors = varparams.maxneighbors

      # determine neighborhood search method
      if varparams.maxneighbors ≠ nothing
        # upper bound in maxneighbors
        maxneighbors > ndata && (maxneighbors = ndata)
        if varparams.neighborhood ≠ nothing
          # local search with a neighborhood
          neigh = varparams.neighborhood

          if neigh isa MetricBall
            bsearcher = KBallSearch(pdata, maxneighbors, neigh)
          else
            searcher  = BallSearch(pdata, neigh)
            bsearcher = BoundedSearch(searcher, maxneighbors)
          end
        else
          # nearest neighbor search with a distance
          distance = varparams.distance
          bsearcher = KNearestSearch(pdata, maxneighbors, metric=distance)
        end
      else
        # use all data points as neighbors
        bsearcher = nothing
      end

      # local inputs
      method = varparams.method
      KC = method == :KernelConvolution ? true : false
      localaniso, localanisohd = (varparams.localaniso, varparams.localanisohd)
      (localanisohd != nothing && KC) && (localanisohd = toqmat(localanisohd))

      # check pars
      okmeth = method in [:MovingWindows, :KernelConvolution]
      @assert okmeth "method must be :MovingWindows or :KernelConvolution"

      oklocal1 = length(localaniso.rotation) == nelements(pdomain)
      oklocal2 = typeof(localaniso) <: LocalAnisotropy
      @assert oklocal1 "number of local anisotropies must match domain points"
      @assert oklocal2 "wrong format of local anisotropies"

      # save preprocessed input
      preproc[var] = (estimator=estimator,
                      minneighbors=minneighbors,
                      maxneighbors=maxneighbors,
                      bsearcher=bsearcher,
                      method=method,
                      localaniso=localaniso,
                      localanisohd=localanisohd)
    end
  end

  preproc
end

function local_solve_approx(problem::EstimationProblem, var::Symbol, preproc)
    # retrieve problem info
    pdata = data(problem)
    pdomain = domain(problem)

    mactypeof = Dict(name(v) => mactype(v) for v in variables(problem))

    # unpack preprocessed parameters
    estimator, minneighbors, maxneighbors, bsearcher, method, localaniso, hdlocalaniso = preproc[var]
    KC = method == :KernelConvolution ? true : false

    # if KC, pass localaniso to hard data
    (KC && hdlocalaniso==nothing) && (hdlocalaniso = grid2hd_qmat(pdata,pdomain,localaniso))

    # determine value type
    V = mactypeof[var]

    # pre-allocate memory for result
    varμ = Vector{V}(undef, nelements(pdomain))
    varσ = Vector{V}(undef, nelements(pdomain))

    # estimation loop
    Threads.@threads for location in traverse(pdomain, LinearPath())
      # pre-allocate memory
      neighbors = Vector{Int}(undef, maxneighbors)

      # centroid of neighborhood center
      pₒ = centroid(pdomain, location)

      # find neighbors with previously estimated values
      nneigh = search!(neighbors, pₒ, bsearcher)

      localpar = (rotation(localaniso,location),magnitude(localaniso,location))
      localestimator = KC ? nothing : mwvario(estimator, localpar)

      # skip location in there are too few neighbors
      if nneigh < minneighbors
        varμ[location] = NaN
        varσ[location] = NaN
      else
        # final set of neighbors
        nview = view(neighbors, 1:nneigh)

        # get neighbors centroid and values
        X = view(pdata, nview)

        # not using block kriging yet, need more tests
        # uₒ = pdomain[location]

        # fit estimator to data and predict mean and variance
        if !KC
          krig = local_fit(localestimator, X)
          μ, σ² = local_predict(krig, var, pₒ)
        else
          ∑neighs = view(hdlocalaniso, nview)
          krig = local_fit(estimator, X, ∑neighs)
          Qx₀ = qmat(localpar...)
          μ, σ² = local_predict(krig, var, pₒ, (Qx₀,∑neighs))
        end

        varμ[location] = μ
        varσ[location] = σ²
      end
    end

    varμ, varσ
end

function local_lhs(estimator::KrigingEstimator, domain,
  localaniso::AbstractVector)

  γ = estimator.γ
  nobs = length(domain)
  ncons = nconstraints(estimator)

  # pre-allocate memory for LHS
  u = first(domain)
  T = Variography.result_type(γ, u, u)
  m = nobs + ncons
  LHS = Matrix{T}(undef, m, m)

  # set variogram/covariance block
  KC = (localaniso[1] == nothing) ? false : true

  if !KC
    Variography.pairwise!(LHS, γ, domain)
    if isstationary(γ)
      for j=1:nobs, i=1:nobs
        @inbounds LHS[i,j] = sill(γ) - LHS[i,j]
      end
    end
  else
    kcfill!(LHS, γ, domain, localaniso)
  end

  # set blocks of constraints
  set_constraints_lhs!(estimator, LHS, domain)

  LHS
end

function set_local_rhs!(estimator::FittedKriging, pₒ,
  localaniso::Tuple)

  γ = estimator.estimator.γ
  X = domain(estimator.state.data)
  RHS = estimator.state.RHS

  KC = localaniso[1]==nothing ? false : true

  # RHS variogram/covariance
  if !KC
    @inbounds for j in 1:nelements(X)
      xj = centroid(X, j)
      RHS[j] = isstationary(γ) ? sill(γ) - γ(xj, pₒ) : γ(xj, pₒ)
    end
  else
    @inbounds for j in 1:nelements(X)
      xj = centroid(X, j)
      RHS[j] = kccov(γ, pₒ, xj, localaniso[1], localaniso[2][j])
    end
  end

  set_constraints_rhs!(estimator, pₒ)
end

function local_fit(estimator::KrigingEstimator, data,
  localaniso::AbstractVector=[nothing])

  # build Kriging system
  LHS = local_lhs(estimator, domain(data), localaniso)
  RHS = Vector{eltype(LHS)}(undef, size(LHS,1))

  # factorize LHS
  FLHS = factorize(estimator, LHS)

  # record Kriging state
  state = KrigingState(data, FLHS, RHS)

  # return fitted estimator
  FittedKriging(estimator, state)
end

function local_weights(estimator::FittedKriging, pₒ,
  localaniso::Tuple)

  nobs = nelements(estimator.state.data)

  set_local_rhs!(estimator, pₒ, localaniso)

  # solve Kriging system
  x = estimator.state.LHS \ estimator.state.RHS

  λ = view(x,1:nobs)
  ν = view(x,nobs+1:length(x))

  KrigingWeights(λ, ν)
end

function local_predict(estimator::FittedKriging, var, pₒ, localaniso::Tuple=(nothing,))
  data = estimator.state.data
  combine(estimator, local_weights(estimator, pₒ, localaniso), data[var])
end