text:
using DSGE, Test, ModelConstructors

# Initialize model object
m = Model990(testing = true)
m <= Setting(:date_forecast_start, quartertodate("2015-Q4"))
m <= Setting(:use_population_forecast, true)

# Forecast under historical rule
paras = [α.value for α in m.parameters]
df = load_data(m, verbose = :none)

output_vars = [:histpseudo, :forecastpseudo, :irfpseudo]
out_hist = DSGE.forecast_one_draw(m, :mode, :none, output_vars, paras, df, verbose = :none)

# Getting indices
eq_mp    = m.equilibrium_conditions[:eq_mp]
eq_other = setdiff(1:n_equilibrium_conditions(m), eq_mp)

# Assign alternative policy
m <= Setting(:alternative_policy, DSGE.taylor93())

# Compare equilibrium conditions under historical and alternative rules
Γ0_hist, Γ1_hist, ~ = eqcond(m)
Γ0_alt93,  Γ1_alt93,  ~ = alternative_policy(m).eqcond(m)

# Repeat for Taylor99
m <= Setting(:alternative_policy, DSGE.taylor99())
Γ0_alt99,  Γ1_alt99,  ~ = alternative_policy(m).eqcond(m)

## First, check if coefficients in matrices correct
### Taylor 93
endo = m.endogenous_states
@testset "Check Taylor93 rule has the right coefficients." begin
    @test Γ0_alt93[eq_mp, endo[:R_t]] == 1
    @test Γ0_alt93[eq_mp, endo[:π_a_t]] == -1.5/4
    @test Γ0_alt93[eq_mp, endo[:y_t]]  == -0.5/4
    @test Γ0_alt93[eq_mp, endo[:y_f_t]] == 0.5/4
end

### Taylor 99
@testset "Check Taylor99 rule has the right coefficients." begin
    @test Γ0_alt99[eq_mp, endo[:R_t]] == 1
    @test Γ0_alt99[eq_mp, endo[:π_a_t]] == -1.5/4
    @test Γ0_alt99[eq_mp, endo[:y_t]]  == -1/4
    @test Γ0_alt99[eq_mp, endo[:y_f_t]] == 1/4
end

# Testing
@testset "Compare AltPolicy eqconds under hist and alt rules" begin
    @test Γ0_hist[eq_other, :] ≈ Γ0_alt93[eq_other, :]
    @test Γ1_hist[eq_other, :] ≈ Γ1_alt93[eq_other, :]
    @test !(Γ0_hist[eq_mp, :]  ≈ Γ0_alt93[eq_mp, :])
    @test !(Γ1_hist[eq_mp, :]  ≈ Γ1_alt93[eq_mp, :])

    @test Γ0_hist[eq_other, :] ≈ Γ0_alt99[eq_other, :]
    @test Γ1_hist[eq_other, :] ≈ Γ1_alt99[eq_other, :]
    @test !(Γ0_hist[eq_mp, :]  ≈ Γ0_alt99[eq_mp, :])
    @test !(Γ1_hist[eq_mp, :]  ≈ Γ1_alt99[eq_mp, :])
end

# Test error thrown if trying to run shockdec products
@testset "Check AltPolicy error thrown with shockdec products" begin
    @test_throws ErrorException forecast_one(m, :mode, :none, [:shockdecobs])
    @test_throws ErrorException forecast_one(m, :mode, :none, [:dettrendobs])
    @test_throws ErrorException forecast_one(m, :mode, :none, [:trendobs])
end

# Forecast under Taylor 93 and Taylor 99
m <= Setting(:alternative_policy, DSGE.taylor93())
out_alt93 = DSGE.forecast_one_draw(m, :mode, :none, output_vars, paras, df, verbose = :none)
m <= Setting(:alternative_policy, DSGE.taylor99())
out_alt99 = DSGE.forecast_one_draw(m, :mode, :none, output_vars, paras, df, verbose = :none)

@testset "Check forecast under Taylor 93" begin
    @test out_hist[:histpseudo] ≈ out_alt93[:histpseudo]
    @test !(out_hist[:forecastpseudo] ≈ out_alt93[:forecastpseudo])
end
@testset "Check forecast under Taylor 99" begin
    @test out_hist[:histpseudo] ≈ out_alt99[:histpseudo]
    @test !(out_hist[:forecastpseudo] ≈ out_alt99[:forecastpseudo])
end

####################
# New Policy Rules #
####################

m = Model1002(; testing = true, custom_settings = Dict{Symbol, Setting}(:add_pgap => Setting(:add_pgap, true)))
my = Model1002(; testing = true, custom_settings = Dict{Symbol, Setting}(:add_pgap => Setting(:add_pgap, true),
                                                                         :add_ygap => Setting(:add_ygap, true)))
m <= Setting(:date_forecast_start, quartertodate("2015-Q4"))
m <= Setting(:use_population_forecast, true)
my <= Setting(:date_forecast_start, quartertodate("2015-Q4"))
my <= Setting(:use_population_forecast, true)

# Getting indices
eq_mp    = m.equilibrium_conditions[:eq_mp]
eq_other = setdiff(1:n_equilibrium_conditions(m), eq_mp)
eq_mpy    = my.equilibrium_conditions[:eq_mp]
eq_othery = setdiff(1:n_equilibrium_conditions(my), eq_mpy)
eq_otherymp = 1:n_equilibrium_conditions(my)

# Compare equilibrium conditions under historical and alternative rules
Γ0_hist, Γ1_hist, ~ = eqcond(m)
Γ0_histy, Γ1_histy, ~ = eqcond(my)

# Matrices for alt_inflation
m <= Setting(:alternative_policy, DSGE.alt_inflation())
Γ0_alt_inf,  Γ1_alt_inf,  ~ = alternative_policy(m).eqcond(m)

# Repeat for zero_rate
m <= Setting(:alternative_policy, DSGE.zero_rate())
Γ0_alt_zero,  Γ1_alt_zero,  ~ = alternative_policy(m).eqcond(m)

# Repeat for rw
my <= Setting(:alternative_policy, DSGE.rw())
Γ0_alt_rw,  Γ1_alt_rw,  ~ = alternative_policy(my).eqcond(my)

# Repeat for rw_zero_rate
my <= Setting(:alternative_policy, DSGE.rw_zero_rate())
Γ0_alt_rw_zero,  Γ1_alt_rw_zero,  ~ = alternative_policy(my).eqcond(my)

# Repeat for smooth_ait_gdp
my <= Setting(:alternative_policy, DSGE.smooth_ait_gdp())
Γ0_alt_smooth,  Γ1_alt_smooth,  ~ = alternative_policy(my).eqcond(my)

# Repeat for smooth_ait_gdp_alt
my <= Setting(:alternative_policy, DSGE.smooth_ait_gdp_alt())

Γ0_alt_smooth_alt,  Γ1_alt_smooth_alt,  ~ = alternative_policy(my).eqcond(my)

# Repeat for ait
my <= Setting(:alternative_policy, DSGE.ait())
my <= Setting(:pgap_type, :ait)
my <= Setting(:pgap_value, 0.)
Γ0_alt_ait,  Γ1_alt_ait,  ~ = alternative_policy(my).eqcond(my)

# Repeat for ngdp
my <= Setting(:alternative_policy, DSGE.ngdp())
my <= Setting(:pgap_type, :ngdp)
my <= Setting(:pgap_value, 12.)
Γ0_alt_ngdp,  Γ1_alt_ngdp,  ~ = alternative_policy(my).eqcond(my)

# Testing
@testset "Compare AltPolicy eqconds under hist and alt rules" begin
    @test Γ0_hist[eq_other, :] ≈ Γ0_alt_inf[eq_other, :]
    @test Γ1_hist[eq_other, :] ≈ Γ1_alt_inf[eq_other, :]
    @test !(Γ0_hist[eq_mp, :]  ≈ Γ0_alt_inf[eq_mp, :])
    @test !(Γ1_hist[eq_mp, :]  ≈ Γ1_alt_inf[eq_mp, :])

    @test Γ0_histy[eq_othery, eq_otherymp] ≈ Γ0_alt_rw[eq_othery, eq_otherymp]
    @test Γ1_histy[eq_othery, eq_otherymp] ≈ Γ1_alt_rw[eq_othery, eq_otherymp]
    @test !(Γ0_histy[eq_mp, eq_other]  ≈ Γ0_alt_rw[eq_mp, eq_othery])
    @test !(Γ1_histy[eq_mp, eq_other]  ≈ Γ1_alt_rw[eq_mp, eq_othery])

    @test Γ0_histy[eq_othery, eq_otherymp] ≈ Γ0_alt_rw_zero[eq_othery, eq_otherymp]
    @test Γ1_histy[eq_othery, eq_otherymp] ≈ Γ1_alt_rw_zero[eq_othery, eq_otherymp]
    @test !(Γ0_histy[eq_mp, eq_other]  ≈ Γ0_alt_rw_zero[eq_mp, eq_othery])
    @test !(Γ1_histy[eq_mp, eq_other]  ≈ Γ1_alt_rw_zero[eq_mp, eq_othery])

    @test Γ0_hist[eq_other, eq_other] ≈ Γ0_alt_zero[eq_other, eq_other]
    @test Γ1_hist[eq_other, eq_other] ≈ Γ1_alt_zero[eq_other, eq_other]
    @test !(Γ0_hist[eq_mp, :]  ≈ Γ0_alt_zero[eq_mp, :])
    @test !(Γ1_hist[eq_mp, :]  ≈ Γ1_alt_zero[eq_mp, :])

    @test Γ0_histy[eq_othery, eq_otherymp] ≈ Γ0_alt_smooth[eq_othery, eq_otherymp]
    @test Γ1_histy[eq_othery, eq_otherymp] ≈ Γ1_alt_smooth[eq_othery, eq_otherymp]
    @test !(Γ0_histy[eq_mp, eq_other]  ≈ Γ0_alt_smooth[eq_mp, eq_othery])
    @test !(Γ1_histy[eq_mp, eq_other]  ≈ Γ1_alt_smooth[eq_mp, eq_othery])

    @test Γ0_histy[eq_othery, eq_otherymp] ≈ Γ0_alt_smooth_alt[eq_othery, eq_otherymp]
    @test Γ1_histy[eq_othery, eq_otherymp] ≈ Γ1_alt_smooth_alt[eq_othery, eq_otherymp]
    @test !(Γ0_histy[eq_mp, eq_other]  ≈ Γ0_alt_smooth_alt[eq_mp, eq_other])
    @test !(Γ1_histy[eq_mp, eq_other]  ≈ Γ1_alt_smooth_alt[eq_mp, eq_other])

    @test Γ0_histy[eq_othery, eq_otherymp] ≈ Γ0_alt_ait[eq_othery, eq_otherymp]
    @test Γ1_histy[eq_othery, eq_otherymp] ≈ Γ1_alt_ait[eq_othery, eq_otherymp]
    @test !(Γ0_histy[eq_mp, eq_other]  ≈ Γ0_alt_ait[eq_mp, eq_othery])
    @test !(Γ1_histy[eq_mp, eq_other]  ≈ Γ1_alt_ait[eq_mp, eq_othery])

    @test Γ0_histy[eq_othery, eq_otherymp] ≈ Γ0_alt_ngdp[eq_othery, eq_otherymp]
    @test Γ1_histy[eq_othery, eq_otherymp] ≈ Γ1_alt_ngdp[eq_othery, eq_otherymp]
    @test !(Γ0_histy[eq_mp, eq_other]  ≈ Γ0_alt_ngdp[eq_mp, eq_othery])
    @test !(Γ1_histy[eq_mp, eq_other]  ≈ Γ1_alt_ngdp[eq_mp, eq_othery])
end

# Checking if setting variables works
## These values are random and not actually good choices
m <= Setting(:ait_Thalf, 14.2)
m <= Setting(:gdp_Thalf, 5.6)
m <= Setting(:smooth_ait_gdp_ρ_smooth, 0.23)
m <= Setting(:smooth_ait_gdp_φ_π, 0.36)
m <= Setting(:smooth_ait_gdp_φ_y, 0.56)
m <= Setting(:smooth_ait_gdp_alt_ρ_smooth, 0.76)
m <= Setting(:smooth_ait_gdp_alt_φ_π, 0.11)
m <= Setting(:smooth_ait_gdp_alt_φ_y, 0.89)
m <= Setting(:rw_ρ_smooth, 0.45)
m <= Setting(:rw_φ_π, 0.83)
m <= Setting(:rw_φ_y, 0.94)
m <= Setting(:ρ_rw, 0.63)

@testset "Checking Settings assignments" begin
    @test get_setting(m, :ait_Thalf) == 14.2
    @test get_setting(m, :gdp_Thalf) == 5.6
    @test get_setting(m, :smooth_ait_gdp_ρ_smooth) == 0.23
    @test get_setting(m, :smooth_ait_gdp_φ_π) == 0.36
    @test get_setting(m, :smooth_ait_gdp_φ_y) == 0.56
    @test get_setting(m, :smooth_ait_gdp_alt_ρ_smooth) == 0.76
    @test get_setting(m, :smooth_ait_gdp_alt_φ_π) == 0.11
    @test get_setting(m, :smooth_ait_gdp_alt_φ_y) == 0.89
    @test get_setting(m, :rw_ρ_smooth) == 0.45
    @test get_setting(m, :rw_φ_π) == 0.83
    @test get_setting(m, :rw_φ_y) == 0.94
    @test get_setting(m, :ρ_rw) == 0.63
end


text:
using BinDeps
using Compat

@BinDeps.setup

# General settings.
so = "so"
all_load = "--whole-archive"
noall_load = "--no-whole-archive"

@static if Sys.isapple()
  so = "dylib"
  all_load = "-all_load"
  noall_load = "-noall_load"
  using Homebrew
end

hsl_modules = @compat Dict()

here = dirname(@__FILE__)

blas_and_metis_built = false

function build_blas_and_metis()
  global blas_and_metis_built
  if !blas_and_metis_built
    @info "building blas"
    include(joinpath(here, "build_blas.jl"))
    @info "building metis"
    include(joinpath(here, "build_metis4.jl"))
    blas_and_metis_built = true
  end
end

const hsl_ma97_version = "2.4.0"
const hsl_ma97_sha256 = "b91552164311add95f7228d1397a747611f08ffdc86a100df58ddfcedfdc7ca7"
const hsl_ma97_archive = joinpath(here, "downloads", "hsl_ma97-$hsl_ma97_version.tar.gz")

@info "looking for $hsl_ma97_archive"
if isfile(hsl_ma97_archive)
  @info "hsl_ma97 found"
  build_blas_and_metis()
  include("build_hsl_ma97.jl")
else
  @info "hsl_ma97 not found"
end

@eval @BinDeps.install $hsl_modules

ext: jl


text:
module MBWREoS

export MBWREoSComponent, MBWREoSMixture
export name, acentric_factor, molar_mass, carbon_number, pressure
export describe

include("constants.jl")
include("types.jl")
include("interface.jl")
include("dbload.jl")
include("basic_thermo.jl")
include("chempotential.jl")
#include("vt_stability.jl")
#include("vt_flash.jl")
#include("newton.jl")

end # module


text:
#using Pkg
#Pkg.activate("submissions/2021feb15")
#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

using Plots
#include("glom_generic.jl")
=#

#rv_indicators_type = "dcpca0"
#kernel_choice = 3
#starid = 101501
#starid = 10700
#starid = 26965
#starid = 34411

#=
if starid == 101501
    star_rot_rate = 16.28  # days
elseif starid == 10700
    star_rot_rate = 34.0  # days  # wikipedia, how embarassing
elseif starid == 26965
    star_rot_rate = 37.1  # days  #  Saar & Osten (1997) via https://iopscience.iop.org/article/10.3847/1538-3881/aaa896
elseif starid == 34411
    star_rot_rate = 27.9  # days # https://onlinelibrary.wiley.com/doi/full/10.1002/asna.202013790
end

init_timescale_1 = star_rot_rate/4
init_timescale_2 = 1.0/24
tighten_lengthscale_priors = 3
include("glom_generic.jl")

=#

#=
starid_list = [101501, 26965, 34411, 10700]
rv_indicators_type_list = ["dcpca0", "scalpels0", "dcpca3", "scalpels3", "dcpca4", "scalpels4" ]
kernel_choice_list = [3, 5, 4]

#for kernel_choice in kernel_choice_list
kernel_choice = 3
    for starid in starid_list
        if starid == 101501
            star_rot_rate = 16.28  # days
        elseif starid == 10700
            star_rot_rate = 34.0  # days  # wikipedia, how embarassing
        elseif starid == 26965
            star_rot_rate = 37.1  # days  #  Saar & Osten (1997) via https://iopscience.iop.org/article/10.3847/1538-3881/aaa896
        elseif starid == 34411
            star_rot_rate = 27.9  # days # https://onlinelibrary.wiley.com/doi/full/10.1002/asna.202013790
        end
        init_timescale_1 = star_rot_rate/4
        init_timescale_2 = 1.0/24
        if kernel_choice == 4 ||  kernel_choice == 6
            tighten_lengthscale_priors = 3
        else
            tighten_lengthscale_priors = 1
        end

        for rv_indicators_type in rv_indicators_type_list
            #run_glom_rv(starid, rv_indicators_type, kernel_choice, star_rot_rate, init_timescale_1, init_timescale_2, tighten_lengthscale_priors)
=#

#function run_glom_rv(starid, rv_indicators_type, kernel_choice, star_rot_rate, init_timescale_1, init_timescale_2, tighten_lengthscale_priors)
#####################################################################
# CTRL+F "CHANGE"TO FIND PLACES WHERE YOU SHOULD MAKE MODIFICATIONS #
#####################################################################

## Problem setup

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

println("# Working on ", starid, ": ", rv_indicators_type, "  ",  kernel_name)
filename_log = string(starid) * "_pennstate_" * rv_indicators_type * "_glom_" * kernel_name * ".log"
io = open(filename_log, "w")
println(io, "# Working on ", starid, ": ", rv_indicators_type, "  ",  kernel_name)

# CHANGE: the stars rotation rate which is used as the first guess for some GLOM
# hyperparameters and starting point for priors
##### starid = 101501
##### starid = 10700
##### starid = 26965
##### starid = 34411
##### rv_indicators_type = "dcpca0"
##### rv_indicators_type = "dcpca3"
##### rv_indicators_type = "dcpca4"
##### rv_indicators_type = "scalpels0"
##### rv_indicators_type = "scalpels3"
##### rv_indicators_type = "scalpels4"

#=
if starid == 101501
    star_rot_rate = 16.28  # days
elseif starid == 10700
    star_rot_rate = 34.0  # days  # wikipedia, how embarassing
elseif starid == 26965
    star_rot_rate = 37.1  # days  #  Saar & Osten (1997) via https://iopscience.iop.org/article/10.3847/1538-3881/aaa896
elseif starid == 34411
    star_rot_rate = 27.9  # days # https://onlinelibrary.wiley.com/doi/full/10.1002/asna.202013790
end
=#

# importing Yale's 101501 data
#data = CSV.read("examples/101501_activity.csv", DataFrame)
data = CSV.read(joinpath("pennstate",string(starid),string(starid) * "_pennstate_" * rv_indicators_type * "_results.csv"), DataFrame)

# CHANGE: observation times go here
obs_xs = collect(data[!, "Time [MJD]"])
# taking out the mean observation times makes the optimization easier for models
# with periodic parameters
GLOM_RV.remove_mean!(obs_xs)

# CHANGE: rvs and their errors go here
obs_rvs = data[!, "RV_C [m/s]"]
#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"))
obs_rvs_err = data[!, "e_RV_C [m/s]"]

# CHANGE: activity indicators and thier errors go here
# you can actually have as many as you want, but obviously it will take longer
# to fit
obs_indicator1 = data[!, "Score_1 [nondim]"]
obs_indicator1_err = repeat([std(obs_indicator1)], length(obs_indicator1))  # I just put something here
obs_indicator2 = data[!, "Score_2 [nondim]"]
obs_indicator2_err = repeat([std(obs_indicator2)], length(obs_indicator2))  # I just put something here

# removing means as the GP model assumes zero mean
GLOM_RV.remove_mean!(obs_rvs)
GLOM_RV.remove_mean!(obs_indicator1)
GLOM_RV.remove_mean!(obs_indicator2)

# CHANGE: change these lines if you add more than 2 indicators
# this takes the data and riffles it together so it takes the form
# [rv_1, ind1_1, ind2_1, rv_2, ind1_2, ind2_2, ...]
n_out = 3  # number of indicators + 1
obs_ys = collect(Iterators.flatten(zip(obs_rvs, obs_indicator1, obs_indicator2)))
obs_noise = collect(Iterators.flatten(zip(obs_rvs_err, obs_indicator1_err, obs_indicator2_err)))

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

# CHANGE: consider changing a0 (the GLOM coefficients that are used, see
# commented lines below)
# If all a's active:
problem_definition = GLOM.GLO(kernel_function, num_kernel_hyperparameters, n_dif, n_out, obs_xs, copy(obs_ys); noise=copy(obs_noise), a0=[[1. 1 1];[1 1 1];[1 1 1]])
#problem_definition_inflated = GLOM.GLO(kernel_function, num_kernel_hyperparameters, n_dif, n_out, obs_xs, copy(obs_ys); noise=copy(obs_noise), a0=[[1. 1 1];[1 1 1];[1 1 1]])
# problem_definition = GLOM.GLO(kernel_function, num_kernel_hyperparameters, n_dif, n_out, obs_xs, copy(obs_ys); noise=copy(obs_noise), a0=[[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 problem_definition.normals
GLOM.normalize_problem_definition!(problem_definition)
#=
GLOM.normalize_problem_definition!(problem_definition_inflated)
for i in 1:problem_definition_inflated.n_out
  inds = (problem_definition_inflated.n_out:length(problem_definition_inflated.y_obs))
  problem_definition_inflated.y_obs[inds] /= problem_definition_inflated.normals[i]
  problem_definition_inflated.noise[inds] /= problem_definition_inflated.normals[i]
end
=#

# CHANGE: Setting initial fit values
initial_total_hyperparameters = collect(Iterators.flatten(problem_definition.a0))
#initial_hypers = [[star_rot_rate], [star_rot_rate], [star_rot_rate], [star_rot_rate, 2 * star_rot_rate, 1], [star_rot_rate, 2 * star_rot_rate, 1], [star_rot_rate, 2 * star_rot_rate, 1]]
##### init_timescale_1 = star_rot_rate
##### init_timescale_2 = 1.0/24
#initial_hypers = [[star_rot_rate], [star_rot_rate], [star_rot_rate], [star_rot_rate, 2 * star_rot_rate, 1], [init_timescale_1, init_timescale_2, 1], [star_rot_rate, 2 * star_rot_rate, 1]]
initial_hypers = [[init_timescale_1], [init_timescale_1], [init_timescale_1], [star_rot_rate, 2 * star_rot_rate, 1], [init_timescale_1, init_timescale_2, 1], [init_timescale_1, init_timescale_2, 1]]
append!(initial_total_hyperparameters, initial_hypers[kernel_choice])

initial_total_hyperparameters

## Fitting GLOM Model

# CHANGE: Setting kernel hyperparameter priors and kick function
# kick functions help avoid saddle points
##### tighten_lengthscale_priors = 1
if kernel_name in ["pp", "se", "m52"]
    kernel_hyper_priors(hps::Vector{<:Real}, d::Integer) =
        #GLOM_RV.kernel_hyper_priors_1λ(hps, d, star_rot_rate, star_rot_rate / 2 / tighten_lengthscale_priors)
        GLOM_RV.kernel_hyper_priors_1λ(hps, d, init_timescale_1, init_timescale_1 / 2 / 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, [star_rot_rate, 2 * star_rot_rate, 1.], [star_rot_rate / 2, star_rot_rate / 2, 0.4] ./ 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, [star_rot_rate, 2 * star_rot_rate, 1.], [star_rot_rate / 2, star_rot_rate / 2, 1.] ./ tighten_lengthscale_priors)
    kernel_hyper_priors(hps::Vector{<:Real}, d::Integer) =
       GLOM_RV.kernel_hyper_priors_2λ(hps, d, [init_timescale_1, init_timescale_2, 1.], [init_timescale_1 / 2, init_timescale_2 / 2, 1.] ./ 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

fit1_total_hyperparameters, result = GLOM_RV.fit_GLOM(problem_definition, initial_total_hyperparameters, kernel_hyper_priors, add_kick!)
# fit_GLOM returns a vector of num_kernel_hyperparameters gp hyperparameters
# followed by the GLOM coefficients and the Optim result object
workspace = GLOM.nlogL_matrix_workspace(problem_definition, fit1_total_hyperparameters)

println("# GP parameters: ", fit1_total_hyperparameters)
println(io,"# GP parameters: ", fit1_total_hyperparameters)

## Plotting initial results

plot_xs = collect(LinRange(obs_xs[1]-10, obs_xs[end]+10, 300))
post, post_err, post_obs = GLOM_RV.GLOM_posteriors(problem_definition, plot_xs, fit1_total_hyperparameters)
GLOM_rvs_at_plot_xs, GLOM_ind1_at_plot_xs, GLOM_ind2_at_plot_xs = post
GLOM_rvs_err_at_plot_xs, GLOM_ind1_err_at_plot_xs, GLOM_ind2_err_at_plot_xs = post_err
GLOM_rvs_at_obs_xs, GLOM_ind1_at_obs_xs, GLOM_ind2_at_obs_xs = post_obs

activity_rvs = GLOM_rvs_at_obs_xs  # the best guess for activity RVs
clean_rvs = obs_rvs - activity_rvs  # the best guess for RVs without activity
#println(" std(obs_rvs) = ",std(obs_rvs), " std(activity_rvs) = ", std(activity_rvs), " std(clean_rvs) = ", std(clean_rvs), " median(sigma_gp) = ", median(GLOM_rvs_err_at_plot_xs) )

#=
using Plots
plt1 = scatter(obs_xs, obs_rvs, yerror=obs_rvs_err, ylabel="RV (m/s)", legend=:none)
plot!(plt1, plot_xs, GLOM_rvs_at_plot_xs, ribbons=GLOM_rvs_err_at_plot_xs, fillalpha=0.3)

plt2 = scatter(obs_xs, obs_indicator1, yerror=obs_indicator1_err, ylabel="Ind 1", legend=:none)
plot!(plt2, plot_xs, GLOM_ind1_at_plot_xs, ribbons=GLOM_ind1_err_at_plot_xs, fillalpha=0.3)

plt3 = scatter(obs_xs, obs_indicator2, yerror=obs_indicator2_err, ylabel="Ind 2", legend=:none)
plot!(plt3, plot_xs, GLOM_ind2_at_plot_xs, ribbons=GLOM_ind2_err_at_plot_xs, fillalpha=0.3)
plt4 = scatter(obs_xs, clean_rvs, yerror=sqrt.(obs_rvs_err.^2 #= .+GLOM_rvs_err_at_obs_xs.^2 =# ), ylabel="RV (m/s)", legend=:none, xlabel="Time (d)")
#plot!(plt4, plot_xs, -GLOM_rvs_at_plot_xs, ribbons=GLOM_rvs_err_at_plot_xs, fillalpha=0.3)
plot(plt1, plt2, plt3, plt4, layout = @layout [a; b; c; d] )
#plot(plt1, plt4, layout = @layout [a; d] )
=#

## Eric's hack to remove RVs from conditioning
#println("# Inflating RV errors and reconditioning")
problem_definition.noise[1:n_out:end] .*= 10000
#GLOM_rvs_at_plot_xs, GLOM_ind1_at_plot_xs, GLOM_ind2_at_plot_xs,
#    GLOM_rvs_err_at_plot_xs, GLOM_ind1_err_at_plot_xs, GLOM_ind2_err_at_plot_xs,
#    GLOM_rvs_at_obs_xs, GLOM_ind1_at_obs_xs, GLOM_ind2_at_obs_xs =
#    posts(problem_definition, fit1_total_hyperparameters)
#after_model_std = std(GLOM_rvs_at_obs_xs - obs_rvs)

post_inflated, post_err_inflated, post_obs_inflated = GLOM_RV.GLOM_posteriors(problem_definition, obs_xs, fit1_total_hyperparameters)
GLOM_rvs_at_plot_xs_inflated, GLOM_ind1_at_plot_xs_inflated, GLOM_ind2_at_plot_xs_inflated = post_inflated
GLOM_rvs_err_at_plot_xs_inflated, GLOM_ind1_err_at_plot_xs_inflated, GLOM_ind2_err_at_plot_xs_inflated = post_err_inflated
GLOM_rvs_at_obs_xs_inflated, GLOM_ind1_at_obs_xs_inflated, GLOM_ind2_at_obs_xs_inflated = post_obs_inflated
GLOM_rvs_err_at_obs_xs_inflated = GLOM_rvs_err_at_plot_xs_inflated

activity_rvs_inflated = GLOM_rvs_at_obs_xs_inflated  # the best guess for activity RVs
clean_rvs_inflated = obs_rvs - activity_rvs_inflated  # the best guess for RVs without activity
println("# Summary: median(obs_rvs_err) = ", median(obs_rvs_err), " std(obs_rvs) = ",std(obs_rvs), " std(activity_rvs) = ", std(activity_rvs_inflated), " std(clean_rvs) = ", std(clean_rvs_inflated), " std(sigma_obs_rvs+activity_rvs) = ", median(sqrt.(obs_rvs_err.^2 .+GLOM_rvs_err_at_obs_xs_inflated.^2)) )
println(io, "# Summary: median(obs_rvs_err) = ", median(obs_rvs_err), " std(obs_rvs) = ",std(obs_rvs), " std(activity_rvs) = ", std(activity_rvs_inflated), " std(clean_rvs) = ", std(clean_rvs_inflated), " std(sigma_obs_rvs+activity_rvs) = ", median(sqrt.(obs_rvs_err.^2 .+GLOM_rvs_err_at_obs_xs_inflated.^2)) )

post_inflated, post_err_inflated, post_obs_inflated = GLOM_RV.GLOM_posteriors(problem_definition, plot_xs, fit1_total_hyperparameters)
GLOM_rvs_at_plot_xs_inflated, GLOM_ind1_at_plot_xs_inflated, GLOM_ind2_at_plot_xs_inflated = post_inflated
GLOM_rvs_err_at_plot_xs_inflated, GLOM_ind1_err_at_plot_xs_inflated, GLOM_ind2_err_at_plot_xs_inflated = post_err_inflated
GLOM_rvs_at_obs_xs_inflated, GLOM_ind1_at_obs_xs_inflated, GLOM_ind2_at_obs_xs_inflated = post_obs_inflated

#------------------
post_inflated, post_err_inflated, post_obs_inflated = GLOM_RV.GLOM_posteriors(problem_definition, obs_xs, fit1_total_hyperparameters)
GLOM_rvs_at_plot_xs_inflated, GLOM_ind1_at_plot_xs_inflated, GLOM_ind2_at_plot_xs_inflated, GLOM_ind3_at_plot_xs_inflated, GLOM_ind4_at_plot_xs_inflated = post_inflated
GLOM_rvs_err_at_plot_xs_inflated, GLOM_ind1_err_at_plot_xs_inflated, GLOM_ind2_err_at_plot_xs_inflated, GLOM_ind3_err_at_plot_xs_inflated, GLOM_ind4_err_at_plot_xs_inflated = post_err_inflated
GLOM_rvs_at_obs_xs_inflated, GLOM_ind1_at_obs_xs_inflated, GLOM_ind2_at_obs_xs_inflated, GLOM_ind3_at_obs_xs_inflated, GLOM_ind4_at_obs_xs_inflated = post_obs_inflated
GLOM_rvs_err_at_obs_xs_inflated = GLOM_rvs_err_at_plot_xs_inflated

activity_rvs_inflated = GLOM_rvs_at_obs_xs_inflated  # the best guess for activity RVs
clean_rvs_inflated = obs_rvs - activity_rvs_inflated  # the best guess for RVs without activity
println("# Summary: median(obs_rvs_err) = ", median(obs_rvs_err), " std(obs_rvs) = ",std(obs_rvs), " std(activity_rvs) = ", std(activity_rvs_inflated), " std(clean_rvs) = ", std(clean_rvs_inflated), " std(sigma_obs_rvs+activity_rvs) = ", median(sqrt.(obs_rvs_err.^2 .+GLOM_rvs_err_at_obs_xs_inflated.^2)) )

post_inflated, post_err_inflated, post_obs_inflated = GLOM_RV.GLOM_posteriors(problem_definition, plot_xs, fit1_total_hyperparameters)
GLOM_rvs_at_plot_xs_inflated, GLOM_ind1_at_plot_xs_inflated, GLOM_ind2_at_plot_xs_inflated, GLOM_ind3_at_plot_xs_inflated, GLOM_ind4_at_plot_xs_inflated = post_inflated
GLOM_rvs_err_at_plot_xs_inflated, GLOM_ind1_err_at_plot_xs_inflated, GLOM_ind2_err_at_plot_xs_inflated, GLOM_ind3_err_at_plot_xs_inflated, GLOM_ind4_err_at_plot_xs_inflated = post_err_inflated
GLOM_rvs_at_obs_xs_inflated, GLOM_ind1_at_obs_xs_inflated, GLOM_ind2_at_obs_xs_inflated, GLOM_ind3_at_obs_xs_inflated, GLOM_ind4_at_obs_xs_inflated = post_obs_inflated

#----
# #=
#using Plots
plt1 = scatter(obs_xs, obs_rvs, yerror=obs_rvs_err, ylabel="RV (m/s)", legend=:none)
plot!(plt1, plot_xs, GLOM_rvs_at_plot_xs_inflated, ribbons=GLOM_rvs_err_at_plot_xs_inflated, fillalpha=0.3)

plt2 = scatter(obs_xs, obs_indicator1, yerror=obs_indicator1_err, ylabel="Ind 1", legend=:none)
plot!(plt2, plot_xs, GLOM_ind1_at_plot_xs_inflated, ribbons=GLOM_ind1_err_at_plot_xs_inflated, fillalpha=0.3)

plt3 = scatter(obs_xs, obs_indicator2, yerror=obs_indicator2_err, ylabel="Ind 2", legend=:none)
plot!(plt3, plot_xs, GLOM_ind2_at_plot_xs_inflated, ribbons=GLOM_ind2_err_at_plot_xs_inflated, fillalpha=0.3)

plt4 = scatter(obs_xs, obs_indicator3, yerror=obs_indicator3_err, ylabel="Ind 3", legend=:none)
plot!(plt4, plot_xs, GLOM_ind3_at_plot_xs_inflated, ribbons=GLOM_ind3_err_at_plot_xs_inflated, fillalpha=0.3)

plt5 = scatter(obs_xs, obs_indicator4, yerror=obs_indicator4_err, ylabel="Ind 4", legend=:none)
plot!(plt5, plot_xs, GLOM_ind4_at_plot_xs_inflated, ribbons=GLOM_ind4_err_at_plot_xs_inflated, fillalpha=0.3)

plt6 = scatter(obs_xs, clean_rvs_inflated, yerror=sqrt.(obs_rvs_err.^2 .+GLOM_rvs_err_at_obs_xs_inflated.^2 ), ylabel="RV (m/s)", legend=:none, xlabel="Time (d)")
plot(plt1, plt2, plt3, plt4, plt5, plt6, layout = @layout [a; b; c; d; e; f] )
filename_fig = joinpath("pennstate",string(starid),string(starid) * "_pennstate_" * rv_indicators_type * "_glom_" * kernel_name * "_results.png")
savefig(filename_fig)
# =#

#names(data)
data_out = deepcopy(data)

if contains(rv_indicators_type,"dcpca")
    data_out[!,"RV_C [m/s]"] = clean_rvs_inflated
    data_out[!,"e_RV_C [m/s]"] = sqrt.(obs_rvs_err.^2 .+GLOM_rvs_err_at_obs_xs_inflated.^2)
    data_out[!,"RV_A [m/s]"] = activity_rvs_inflated
    data_out[!,"e_RV_A [m/s]"] = GLOM_rvs_err_at_obs_xs_inflated
elseif contains(rv_indicators_type,"scalpel")
    data_out[!,"RV_C [m/s]"] = clean_rvs_inflated
    data_out[!,"e_RV_C [m/s]"] = sqrt.(obs_rvs_err.^2 .+GLOM_rvs_err_at_obs_xs_inflated.^2)
    data_out[!,"RV_A [m/s]"] .+= activity_rvs_inflated
    data_out[!,"e_RV_A [m/s]"] .= GLOM_rvs_err_at_obs_xs_inflated
end
for i in 1:2
   #output_df[!,"Score_" * string(i) * " [nondim]"] = vec(scores[:,i])
   data_out[!,"e_Score_" * string(i) * " [nondim]"] .= std(data[!,"Score_" * string(i) * " [nondim]"])
end

filename_out = joinpath("pennstate",string(starid),string(starid) * "_pennstate_" * rv_indicators_type * "_glom_" * kernel_name * "_results.csv")
CSV.write(filename_out, data_out)

println(io,"# Done")
close(io)
println("# Done")

## Coefficient exploration

#=
# Could use this to iterate through all of the possible combinations of GLOM
# coefficients
possible_a0s = Matrix[]
GLOM_RV.valid_a0s!(possible_a0s, zeros(n_out, n_dif))
append!(possible_a0s, [ones(n_out, n_dif)])
reverse!(possible_a0s)  # the expensive, filled out possibilities go first
nℓs = Float64[]
all_fit_total_hyperparameters = Vector[]
problem_definitions = GLOM.GLO[]
t0 = Libc.time()
for i in 1:length(possible_a0s)
    a0 = possible_a0s[i]
    problem_definition = GLOM.GLO(kernel_function, num_kernel_hyperparameters, n_dif, n_out, obs_xs, copy(obs_ys); noise=copy(obs_noise), a0=a0)
    GLOM.normalize_problem_definition!(problem_definition)
    initial_total_hyperparameters = collect(Iterators.flatten(problem_definition.a0))
    append!(initial_total_hyperparameters, initial_hypers[kernel_choice])
    fit_total_hyperparameters, result = GLOM_RV.fit_GLOM(problem_definition, initial_total_hyperparameters, kernel_hyper_priors, add_kick!; print_stuff=false)
    append!(problem_definitions, [problem_definition])
    append!(nℓs, result.minimum)
    append!(all_fit_total_hyperparameters, [fit_total_hyperparameters])
    println("\nDone with $(round(100 * i / length(possible_a0s); digits=2))% of a0 possibilities")
    t = Libc.time() - t0
    println("t: $(Int(round(t)))s t_left?: $(round(t * (length(possible_a0s) - i) / i / 60; digits=1)) mins")
end

best_fits = sortperm(nℓs)

plot_xs = collect(LinRange(obs_xs[1]-10, obs_xs[end]+10, 300))
post, post_err, post_obs = GLOM_RV.GLOM_posteriors(problem_definitions[best_fits[1]], plot_xs, all_fit_total_hyperparameters[best_fits[1]])
GLOM_rvs_at_plot_xs, GLOM_ind1_at_plot_xs, GLOM_ind2_at_plot_xs = post
GLOM_rvs_err_at_plot_xs, GLOM_ind1_err_at_plot_xs, GLOM_ind2_err_at_plot_xs = post_err
GLOM_rvs_at_obs_xs, GLOM_ind1_at_obs_xs, GLOM_ind2_at_obs_xs = post_obs

using Plots
plt = scatter(obs_xs, obs_rvs, yerror=obs_rvs_err)
plot!(plt, plot_xs, GLOM_rvs_at_plot_xs, ribbons=GLOM_rvs_err_at_plot_xs, fillalpha=0.3)

plt = scatter(obs_xs, obs_indicator1, yerror=obs_indicator1_err)
plot!(plt, plot_xs, GLOM_ind1_at_plot_xs, ribbons=GLOM_ind1_err_at_plot_xs, fillalpha=0.3)

plt = scatter(obs_xs, obs_indicator2, yerror=obs_indicator2_err)
plot!(plt, plot_xs, GLOM_ind2_at_plot_xs, ribbons=GLOM_ind2_err_at_plot_xs, fillalpha=0.3)
=#

## Finding a planet?

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

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

println("starting hyperparameters")
println(initial_total_hyperparameters)
initial_nlogL = GLOM.nlogL_GLOM(problem_definition, 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(problem_definition, 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(problem_definition_rv.time; samples_per_peak=11)
period_grid = 1 ./ reverse(freq_grid)
amount_of_periods = length(period_grid)

Σ_obs = GLOM.Σ_observations(problem_definition, 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(), problem_definition_rv=problem_definition_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(problem_definition_rv, Σ_obs, GLOM_RV.kep_signal_epicyclic(P=P))
    if !fast
        ks = GLOM_RV.fit_kepler(problem_definition_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(problem_definition_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(
            problem_definition_rv.GLO,
            fit1_total_hyperparameters;
            Σ_obs=Σ_obs,
            y_obs=GLOM_RV.remove_kepler(problem_definition_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)

likelihoods = zeros(amount_of_periods)
unnorm_posteriors = zeros(amount_of_periods)

@time if use_distributed
    # takes around 8s for 101501 data with 6 workers and 1200 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 15s for 101501 data and 1200 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")

# using Plots
# plot(ustrip.(period_grid), likelihoods; xaxis=:log, leg=false)
# plot(ustrip.(period_grid), unnorm_posteriors; xaxis=:log, leg=false)

####################################################################################################
# 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]

current_ks = fit_kep_hold_P(best_period; print_stuff=true)
println("before wright fit: ", GLOM_RV.kep_parms_str(current_ks))

#=
plot_kep_xs = collect(LinRange(0, ustrip(best_period), 1000))
# scatter(remainder(problem_definition.x_obs, ustrip(best_period)), ustrip.(problem_definition_rv.rv); yerror=ustrip.(problem_definition_rv.rv_noise), label="data")
plot!(plot_kep_xs, ustrip.(current_ks.(plot_kep_xs.*u"d")); label="kep")
=#
fit2_total_hyperparameters, current_ks = GLOM_RV.fit_GLOM_and_kep!(workspace,
    problem_definition_rv, fit1_total_hyperparameters, kernel_hyper_priors,
    add_kick!, current_ks; avoid_saddle=false)

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

current_ks = GLOM_RV.kep_signal(current_ks)
println("\nbefore full fit: ", GLOM_RV.kep_parms_str(current_ks))
fit3_total_hyperparameters, current_ks = GLOM_RV.fit_GLOM_and_kep!(workspace,
    problem_definition_rv, fit2_total_hyperparameters, kernel_hyper_priors,
    add_kick!, current_ks; avoid_saddle=true)

full_ks = GLOM_RV.kep_signal(current_ks)

# # these should be near 0
# GLOM_RV.test_∇nlogL_kep(problem_definition_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(fit3_total_hyperparameters)
fit_nlogL2 = GLOM.nlogL_GLOM!(workspace, problem_definition, fit3_total_hyperparameters; y_obs=GLOM_RV.remove_kepler(problem_definition_rv, full_ks))
uE2 = -fit_nlogL2 - nlogprior_hyperparameters(fit3_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 #
##################################################################################

fit1_total_hyperparameters_temp, result = GLOM_RV.fit_GLOM(
    problem_definition,
    fit3_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!(workspace, problem_definition, 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
end

##########################
# Evidence approximation #
##########################

# no planet
H1 = (GLOM.∇∇nlogL_GLOM(problem_definition, 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!(workspace, problem_definition_rv, fit3_total_hyperparameters, full_ks; include_kepler_priors=true))
n_hyper = length(GLOM.remove_zeros(fit3_total_hyperparameters))
H2[1:n_hyper, 1:n_hyper] += nlogprior_hyperparameters(GLOM.remove_zeros(fit3_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(problem_definition_rv, workspace.Σ_obs, current_ks; include_priors=true)

=#
#end
#end
#end

#end


text:
@testset "code generation" begin

############################################################################################

@testset "LLVM IR" begin

@testset "basic reflection" begin
    llvm_dummy() = return nothing
    ir = sprint(io->CUDAnative.code_llvm(io, llvm_dummy, (); optimize=false, dump_module=true))

    # module should contain our function + a generic call wrapper
    @test contains(ir, "define void @julia_llvm_dummy")
    @test !contains(ir, "define %jl_value_t* @jlcall_")
    @test ismatch(r"define void @julia_llvm_dummy_.+\(\) #0.+\{", ir)
end

@testset "exceptions" begin
    @eval codegen_exception() = throw(DivideError())
    ir = sprint(io->CUDAnative.code_llvm(io, codegen_exception, ()))

    # exceptions should get lowered to a plain trap...
    @test contains(ir, "llvm.trap")
    # not a jl_throw referencing a jl_value_t representing the exception
    @test !contains(ir, "jl_value_t")
    @test !contains(ir, "jl_throw")
end

@testset "sysimg" begin
    # bug: use a system image function

    @eval function codegen_call_sysimg(a,i)
        Base.pointerset(a, 0, mod1(i,10), 8)
        return nothing
    end

    ir = sprint(io->CUDAnative.code_llvm(io, codegen_call_sysimg, (Ptr{Int},Int)))
    @test !contains(ir, "jlsys_")
end

@testset "child functions" begin
    # we often test using @noinline child functions, so test whether these survive
    # (despite not having side-effects)
    @eval @noinline codegen_child(i) = i+1
    @eval codegen_parent(i) = (codegen_child(i); nothing)

    ir = sprint(io->CUDAnative.code_llvm(io, codegen_parent, (Int,)))
    @test ismatch(r"call .+ @julia_codegen_child_", ir)
end

@testset "JuliaLang/julia#21121" begin
    @eval function codegen_tuple_leak()
        weight_matrix = @cuStaticSharedMem(Float32, (16, 16))
        sync_threads()
        weight_matrix[1, 16] *= 2
        sync_threads()

        return nothing
    end

    ir = sprint(io->CUDAnative.code_llvm(io, codegen_tuple_leak, ()))
    @test !contains(ir, "inttoptr")
end

end

############################################################################################

@testset "PTX assembly" begin

@testset "basic reflection" begin
    @eval ptx_valid_kernel(a) = (a[0] = 1; nothing)
    @eval ptx_invalid_kernel(a) = (a[0] = 1; 1)

    @test code_ptx(DevNull, ptx_valid_kernel, Tuple{CuDeviceArray{Int,1}}) == nothing
    @test code_ptx(DevNull, ptx_invalid_kernel, Tuple{CuDeviceArray{Int,1}}) == nothing
    @test_throws ErrorException code_ptx(DevNull, ptx_invalid_kernel, Tuple{CuDeviceArray{Int,1}}; kernel=true) == nothing
end

@testset "child functions" begin
    # we often test using @noinline child functions, so test whether these survive
    # (despite not having side-effects)
    @eval @noinline ptx_child(i) = i+1
    @eval ptx_parent(i) = (ptx_child(i); nothing)

    asm = sprint(io->code_ptx(io, ptx_parent, (Int64,)))
    @test ismatch(r"call.uni \(retval0\),\s+julia_ptx_child_"m, asm)
end

@testset "entry-point functions" begin
    @eval @noinline ptx_nonentry(i) = i+1
    @eval ptx_entry(i) = (ptx_nonentry(i); nothing)

    asm = sprint(io->code_ptx(io, ptx_entry, (Int64,); kernel=true))
    @test ismatch(r"\.visible \.entry julia_ptx_entry_", asm)
    @test ismatch(r"\.visible \.func .+ julia_ptx_nonentry_", asm)
end

@testset "delayed lookup" begin
    @eval codegen_ref_nonexisting() = nonexisting
    @test_throws ErrorException code_ptx(codegen_ref_nonexisting, ())
end

@testset "generic call" begin
    @eval codegen_call_nonexisting() = nonexisting()
    @test_throws ErrorException code_ptx(codegen_call_nonexisting, ())
end

@testset "idempotency" begin
    # bug: generate code twice for the same kernel (jl_to_ptx wasn't idempotent)

    @eval codegen_idempotency() = return nothing
    code_ptx(DevNull, codegen_idempotency, ())
    code_ptx(DevNull, codegen_idempotency, ())
end

@testset "child function reuse" begin
    # bug: depending on a child function from multiple parents resulted in
    #      the child only being present once

    @eval @noinline codegen_child_reuse_child(i) = i+1
    @eval function codegen_child_reuse_parent1(i)
        codegen_child_reuse_child(i)
        return nothing
    end

    asm = sprint(io->code_ptx(io, codegen_child_reuse_parent1, (Int,)))
    @test ismatch(r".func .+ julia_codegen_child_reuse_child", asm)

    @eval function codegen_child_reuse_parent2(i)
        codegen_child_reuse_child(i+1)
        return nothing
    end

    asm = sprint(io->code_ptx(io, codegen_child_reuse_parent2, (Int,)))
    @test ismatch(r".func .+ julia_codegen_child_reuse_child", asm)
end

@testset "child function reuse bis" begin
    # bug: similar, but slightly different issue as above
    #      in the case of two child functions
    @eval @noinline codegen_child_reuse_bis_child1(i) = i+1
    @eval @noinline codegen_child_reuse_bis_child2(i) = i+2
    @eval function codegen_child_reuse_bis_parent1(i)
        codegen_child_reuse_bis_child1(i) + codegen_child_reuse_bis_child2(i)
        return nothing
    end
    asm = sprint(io->code_ptx(io, codegen_child_reuse_bis_parent1, (Int,)))

    @eval function codegen_child_reuse_bis_parent2(i)
        codegen_child_reuse_bis_child1(i+1) + codegen_child_reuse_bis_child2(i+1)
        return nothing
    end
    asm = sprint(io->code_ptx(io, codegen_child_reuse_bis_parent2, (Int,)))
end

@testset "indirect sysimg function use" begin
    # issue #9: re-using sysimg functions should force recompilation
    #           (host fldmod1->mod1 throwsm, so the PTX code shouldn't contain a throw)

    # FIXME: Int64 because of #49

    @eval function codegen_recompile(out)
        wid, lane = fldmod1(unsafe_load(out), Int64(32))
        unsafe_store!(out, wid)
        return nothing
    end

    asm = sprint(io->code_ptx(io, codegen_recompile, (Ptr{Int64},)))
    @test !contains(asm, "jl_throw")
    @test !contains(asm, "jl_invoke")   # forced recompilation should still not invoke
end

@testset "compile for host after PTX" begin
    # issue #11: re-using host functions after PTX compilation
    @eval @noinline codegen_recompile_bis_child(x) = x+1

    @eval function codegen_recompile_bis_fromhost()
        codegen_recompile_bis_child(10)
    end

    @eval function codegen_recompile_bis_fromptx()
        codegen_recompile_bis_child(10)
        return nothing
    end

    code_ptx(DevNull, codegen_recompile_bis_fromptx, ())
    @test codegen_recompile_bis_fromhost() == 11
end

@testset "LLVM intrinsics" begin
    # issue #13 (a): cannot select trunc
    @eval codegen_issue_13(x) = convert(Int, x)
    code_ptx(DevNull, codegen_issue_13, (Float64,))
end

end

############################################################################################

@testset "SASS" begin

@testset "basic reflection" begin
    @eval sass_valid_kernel(a) = (a[0] = 1; nothing)
    @eval sass_invalid_kernel(a) = (a[0] = 1; 1)

    @test code_sass(DevNull, sass_valid_kernel, Tuple{CuDeviceArray{Int,1}}) == nothing
    @test_throws ErrorException code_sass(DevNull, sass_invalid_kernel, Tuple{CuDeviceArray{Int,1}}) == nothing
end

end

############################################################################################

end


text:
include("PRM.jl")
#using Distributions

gr()

#### HELPER FXN

    function areaRect(r::HyperRectangle)
        w, h = widths(r)
        area = w*h
    end

    function unionAreaRects(r1::HyperRectangle, r2::HyperRectangle)
        area1 = areaRect(r1)
        area2 = areaRect(r2)
        overlap = areaRect(intersect(r1, r2))
        area = area1 + area2 - overlap
        return area
    end
### END

## PARAMETERS
numSamples = 500
connectRadius = 2
param = algT.AlgParameters(numSamples, connectRadius)

targetNumObs = 3
clutterPercentage = 0.15
roomWidth,roomHeight  = 20,20

startstate = Point(4.,4)
goalstate = Point(15.,15)

## INIT
obstacles = Vector{HyperRectangle}()
perimeter = HyperRectangle(Vec(0.,0), Vec(roomWidth,roomHeight))
wallsPerimeter = algfxn.decompRect(perimeter)

walls = Vector{LineSegment}()
for l in wallsPerimeter
    push!(walls, l)
end

# implement function to create clutter
# Let's say on average we want to create 4 obstacles... (otherwise, *most* of
# the time we will product one large rectangle that is falling out of the room

#targetNumObs = 2

roomArea = roomWidth * roomHeight
targetSumObsArea= roomArea * clutterPercentage
sumObsArea = 0

while sumObsArea < targetSumObsArea
    #x,y = rand(Uniform(1, roomWidth),2)
    x,y = rand(1.0:roomWidth,2)

    randWidth, randHeight = rand(Uniform(1, roomWidth/targetNumObs),2)
    #randWidth, randHeight = rand(1:roomWidth/targetNumObs,2)
    protoObstacle = HyperRectangle(Vec(x, y), Vec(randWidth, randHeight)) #Todo
    if contains(perimeter, protoObstacle)
        push!(obstacles, protoObstacle)
        sumObsArea += randWidth*randHeight
    end
end
print("obstacles generated")

#plot!(walls, color =:black)
#plot!(obstacles, fillalpha=0.5)

@show targetSumObsArea
@show sumObsArea

### Run PRM on cluttered room

r = algT.Room(roomWidth,roomHeight,walls,obstacles)
plotfxn.plotRoom(r)

## Run preprocessing
nodeslist, edgeslist = preprocessPRM(r, param)
print("\n ---- pre-processed --- \n")

## Query created RM

pathcost, isPathFound, solPath = queryPRM(startstate, goalstate, nodeslist, edgeslist, obstacles)
print("\n ---- queried ---- \n")

## Plot path found
timestamp = Base.Dates.now()
title = "PRM with # samples=$numSamples, maxDist=$connectRadius, \npathcost = $pathcost,
        timestamp=$timestamp)\n\n"
roadmap = algT.roadmap(startstate, goalstate, nodeslist, edgeslist)

plot = plotfxn.plotPRM(roadmap, solPath, title::String)
gui()

print("\n --- Time --- \n")
@show timestamp
print("\n --- Obstacles --- \n")
@show obstacles
print("\n --- Nodes --- \n")
@show nodeslist
print("\n --- Edges --- \n")
@show edgeslist
print("\n -------- \n")

# todo: save one "representative" figure from each run

@show listCosts
@show listpSucc
sizeplot = (800, 800)

# sanity check
plot()

supTitle="\nPRM with maxDist=$connectRadius, nTrials=$nTrials.
        (time to run:$timeExperiment, timestamp=$timestamp)\n\n"
costTitle= "numsamples vs pathcost\n"

pPRMcost = scatter(nSamples_list, listCosts',
    color=:black,
    title = supTitle * costTitle, ylabel = "euclidean path cost", xlabel = "numsamples",
    yaxis=((0,100), 0:20:100))

successTitle = "\nnumsamples vs pSuccess\n"
pPRMsuccess = scatter(nSamples_list, listpSucc',
    color = :orange, markersize= 6,
    title = successTitle, ylabel = ("P(success)=numSucc/$nTrials trials"), xlabel = "numsamples",
    yaxis=((0, 1.2), 0:0.1:1))

plot(pPRMcost, pPRMsuccess, layout=(2,1), legend=false,
    xaxis=((0, 320), 0:50:300),
    size = sizeplot, sizeplot = (800, 800))


text:
using Griddly

image_path = joinpath(@__DIR__,"..","..","..","resources","images")
shader_path = joinpath(@__DIR__,"..","..","..","resources","shaders")
gdy_path = joinpath(@__DIR__,"..","..","..","resources","games")

gdy_reader = Griddly.GDYReader(image_path,shader_path)
# grid = Griddly.load!(gdy_reader,joinpath(gdy_path,"Single-Player/Mini-Grid/minigrid-doggo.yaml"))
#
# println("Grid is loaded, go to create the game now")
#
# game = Griddly.create_game(grid,Griddly.ASCII)
#
# println("Game is created, we now register a player")
#
# # Create a player
# player1 = Griddly.register_player!(game,"Bob", Griddly.BLOCK_2D)
#
# println("Player is created, now we initialize the  game")
#
# Griddly.init!(game)
#
# println("Game is initialized now we want to play and observe with an ASCII representation")
#
# Griddly.load_level!(grid,2)
# Griddly.reset!(game)
# list_object_ids = Griddly.get_object_ids_list(grid)
# println("list_ids:$(convert(Array{Int32,1},list_object_ids))")
# object1_char = Griddly.get_object_char_from_id(grid,1)
# println("object1_char:$object1_char")
# mapping = Griddly.get_map_object_ids_char(grid)
# println(mapping)
# observation = Griddly.observe(game)
# observation = convert(Array{Int8,3},Griddly.get_data(observation))
# Griddly.pretty_print_ASCII(grid,mapping,observation)
# for j in 1:200
#     dir = rand(0:5)
#
#     reward, done = Griddly.step_player!(player1,"move", [dir])
#
#     observation = Griddly.observe(game)
#     Griddly.pretty_print_ASCII(grid,mapping,convert(Array{Int8,3},Griddly.get_data(observation)))
# end

#--------------Player ASCII--------------
grid = Griddly.load!(gdy_reader,joinpath(gdy_path,"Single-Player/Mini-Grid/minigrid-doggo.yaml"))

println("Grid is loaded, go to create the game now")

game = Griddly.create_game(grid,Griddly.SPRITE_2D)

println("Game is created, we now register a player")

# Create a player
player1 = Griddly.register_player!(game,"Bob", Griddly.ASCII)

println("Player is created, now we initialize the  game")

Griddly.init!(game)

println("Game is initialized now we want to play and observe with an ASCII representation")

Griddly.load_level!(grid,2)
Griddly.reset!(game)
list_object_ids = Griddly.get_object_ids_list(grid)
println("list_ids:$(convert(Array{Int32,1},list_object_ids))")
object1_char = Griddly.get_object_char_from_id(grid,1)
println("object1_char:$object1_char")
mapping = Griddly.get_map_object_ids_char(grid)
println(mapping)

for j in 1:200
    dir = rand(0:5)

    reward, done = Griddly.step_player!(player1,"move", [dir])

    player1_tiles = Griddly.observe(player1)
    Griddly.pretty_print_ASCII(grid,mapping,convert(Array{Int8,3},Griddly.get_data(player1_tiles)))
end

#--------------2PlayersGame------------------")
#
# grid = Griddly.load!(gdy_reader,joinpath(gdy_path,"RTS/basicRTS.yaml"))
#
# println("Grid is loaded, go to create the game now")
#
# game = Griddly.create_game(grid,Griddly.ASCII)
#
# println("Game is created, we now register a player")
#
# # Create a player
# player1 = Griddly.register_player!(game,"Bob", Griddly.BLOCK_2D)
# player2 = Griddly.register_player!(game,"Alice", Griddly.BLOCK_2D)
#
# println("Players are created, now we initialize the  game")
#
# Griddly.init!(game)
#
# # Load a custom string
# level_string = """  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  W  .  .  .  .  .  M  M  M
#                     .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  M  M
#                     .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  M
#                     .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  .  .  .  .  .  .  M  W  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  .  .  .  .  .  .  M  W  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  H1 .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  W  M  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  B1 .  .  .  .  .  .  .  M  W  M  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  .  .  .  .  .  .  M  W  .  .  .  .  .  .  .  W  W  W  W  W  .  .  .  .
#                     .  .  .  .  W  W  W  W  W  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .  .  B2 .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  W  M  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  W  M  .  .  .  .  .  .  .  .  H2 .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
#                     .  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
#                     M  .  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
#                     M  M  .  .  .  .  .  .  W  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
#                     M  M  M  .  .  .  .  .  W  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
# """
# Griddly.load_level_string!(grid,level_string)
# Griddly.reset!(game)
#
# println("The game is ready to be played")
#
# list_object_ids = Griddly.get_object_ids_list(grid)
# println("list_ids:$(convert(Array{Int32,1},list_object_ids))")
# object1_char = Griddly.get_object_char_from_id(grid,1)
# println("object1_char:$object1_char")
# mapping = Griddly.get_map_object_ids_char(grid)
# println(mapping)
# actions_map = Griddly.get_player_actions_dict(grid)
# println(actions_map)
#
# function run_test()
#     player1_x::Int64 = 6
#     player1_y::Int64 = 11
#
#     player2_x::Int64 = 25
#     player2_y::Int64 = 19
#
#     for j in 1:200
#
#         # Alternate between player_1 and player_2 actions
#         action = collect(keys(actions_map))[rand(1:length(actions_map))]
#         dir = rand(1:length(actions_map[action]))
#         println("action:$action,dir:$dir,player_id:$(j%2+1)")
#         if j % 2 == 0
#             action_dir = [player1_x,player1_y,dir]
#             player_1_step_result = Griddly.step_player!(player1,action,action_dir)
#             if action=="move"
#                 if dir==1
#                     player1_x -=1
#                 elseif dir==2
#                     player1_y -=1
#                 elseif dir ==3
#                     player1_x +=1
#                 else
#                     player1_y +=1
#                 end
#             end
#         else
#             action_dir = [player2_x,player2_y,dir]
#             player_2_step_result = Griddly.step_player!(player2,action,action_dir)
#             if action=="move"
#                 if dir==1
#                     player2_x -=1
#                 elseif dir==2
#                     player2_y -=1
#                 elseif dir ==3
#                     player2_x +=1
#                 else
#                     player2_y +=1
#                 end
#             end
#         end
#
#         observation = Griddly.observe(game)
#         Griddly.pretty_print_ASCII(grid,mapping,convert(Array{Int8,3},Griddly.get_data(observation)))
#     end
# end
# run_test()


text:
export Nearest, Previous, Next

abstract type ConstantInterpType end
struct Nearest <: ConstantInterpType end
struct Previous <: ConstantInterpType end
struct Next <: ConstantInterpType end

struct Constant{T<:ConstantInterpType,BC<:Union{Throw{OnGrid},Periodic{OnCell}}} <: DegreeBC{0}
    bc::BC
    function Constant{T, BC}(bc::BC=BC()) where {T<:ConstantInterpType, BC<:Union{Throw{OnGrid},Periodic{OnCell}}}
        new{T, BC}(bc)
    end
end

# Default to Nearest and Throw{OnGrid}
Constant(args...) = Constant{Nearest}(args...)
Constant{T}() where {T<:ConstantInterpType} = Constant{T,Throw{OnGrid}}(Throw(OnGrid()))
Constant{T}(bc::BC) where {T<:ConstantInterpType,BC<:BoundaryCondition} = Constant{T,BC}(bc)
Constant{T}(::Periodic{Nothing}) where {T<:ConstantInterpType} = Constant{T,Periodic{OnCell}}(Periodic(OnCell()))

function Base.show(io::IO, deg::Constant)
    print(io, nameof(typeof(deg)), '{', typeof(deg).parameters[1], '}', '(')
    show(io, deg.bc)
    print(io, ')')
end

function Base.show(io::IO, deg::Constant{T,Throw{OnGrid}}) where {T <: ConstantInterpType}
    print(io, nameof(typeof(deg)), '{', typeof(deg).parameters[1], '}', '(', ')')
end

"""
Constant b-splines are *nearest-neighbor* interpolations, and effectively
return `A[round(Int,x)]` when interpolating.
"""
Constant

function positions(c::Constant{Previous}, ax, x)  # discontinuity occurs at integer locations
    xm = floorbounds(x, ax)
    δx = x - xm
    fast_trunc(Int, xm), δx
end
function positions(c::Constant{Next}, ax, x)  # discontinuity occurs at integer locations
    xm = ceilbounds(x, ax)
    δx = x - xm
    fast_trunc(Int, xm), δx
end
function positions(c::Constant{Nearest}, ax, x)  # discontinuity occurs at half-integer locations
    xm = roundbounds(x, ax)
    δx = x - xm
    fast_trunc(Int, xm), δx
end

function positions(c::Constant{Previous,Periodic{OnCell}}, ax, x)
    # We do not use floorbounds because we do not want to add a half at
    # the lowerbound to round up.
    xm = floor(x)
    δx = x - xm
    modrange(fast_trunc(Int, xm), ax), δx
end
function positions(c::Constant{Next,Periodic{OnCell}}, ax, x)  # discontinuity occurs at integer locations
    xm = ceilbounds(x, ax)
    δx = x - xm
    modrange(fast_trunc(Int, xm), ax), δx
end

value_weights(::Constant, δx) = (1,)
gradient_weights(::Constant, δx) = (0,)
hessian_weights(::Constant, δx) = (0,)

padded_axis(ax::AbstractUnitRange, ::BSpline{<:Constant}) = ax


text:
# using Distributions, BenchmarkTools
#
function generateSPDmatrix(n)
	A = rand(n)
	m = 0.5 * (A * A')
	PDmat = m + (n * Diagonal(ones(n)))
end

function generateRandomVCM(n::Int64, p::Int64, d::Int64, m::Int64)
	# n-by-p design matrix
	X = randn(n, p)

	# p-by-d mean component regression coefficient for each trait
	B = hcat(ones(p, 1), rand(p))

	V = ntuple(x -> zeros(n, n), m)
	for i = 1:m-1
	  copy!(V[i], generateSPDmatrix(n))
	end
	copy!(V[end], Diagonal(ones(n))) # last covarianec matrix is identity

	# a tuple of m d-by-d variance component parameters
	Σ = ntuple(x -> zeros(d, d), m)
	for i in 1:m
	  copy!(Σ[i], generateSPDmatrix(d))
	end
	return(X, B, Σ, V)
end

### First using the TraitSimulation.jl package

function MN_J(X, B, Σ, V; n_reps = 1)
    n, p = size(X*B)
    sim = [zeros(n, p) for i in 1:n_reps]
    for i in 1:n_reps
        sim[i] = rand(MatrixNormal(X*B, V, Σ))
    end
    return(sim)
end

### Now using the Distributions.jl package
function MN_Jm(X, B, Σ, V; n_reps = 1)
    n, p = size(X*B)
    m = length(V)
    sim = [zeros(n, p) for i in 1:n_reps]
    for i in 1:n_reps
        for j in 1:m
            dist = MatrixNormal(X*B, V[j], Σ[j])
            sim[i] += rand(dist)
        end
    end
    return(sim)
end

function CompareWithJulia(X, B, Σ, V; n_reps = 1)
	d = size(B, 2)
	vecVC = [VarianceComponent(Σ[i], V[i]) for i in eachindex(V)]
	LMMtraitobjm = VCMTrait(X*B, vecVC)
	m1trait = simulate(LMMtraitobjm, n_reps)

	m1julia = MN_Jm(X, B, Σ, V; n_reps = n_reps)[:]

	m1traitall = [collect(m1trait[:, i, :])[:] for i in 1:d]

	m1traitlong = hcat(m1traitall...)
	m1julialong = vcat(m1julia...)

	dfwrite = DataFrame(hcat(m1traitlong, m1julialong))
	return(dfwrite)
end


text:
function getExpLevels(;df=DataTables.RNAseq, condition="unstressed", Max_expr=Max_expr,idxName=1)
    prNames = df[!,idxName]
    ex = df[!, condition]
    ex = capVec(ex, Max_expr)
    return NamedTuple{tuple(Symbol.(prNames)...)}(tuple(ex...))
end

function getConditions(;df=DataTables.BoolCond, drop_missing=true)
    if drop_missing
        return dropmissing(df)
    else
        return df
    end
end