text:
# Script for evaluating the GP model

using PyPlot

include(pwd()*"/utilities/featurescaling.jl")
include(pwd()*"/utilities/normplot.jl")
include(pwd()*"/utilities/posteriorinference.jl")

try
  cd("Ricker model")
catch
 warn("Already in the Ricker model folder")
end

include("rickermodel.jl")

try
  cd("..")
catch
 warn("Already in main folder")
end

################################################################################
##     set up  problem                                                        ##
################################################################################

problem = set_up_gp_problem(ploton = false)

# set adaptive updating
#problem.adaptive_update = AMUpdate(eye(3), 2.4/sqrt(3), 1., 0.7, 50)

#problem.adaptive_update = noAdaptation(2.4/sqrt(3)*eye(3))

# or, use AM gen alg for adaptive updating
#problem.adaptive_update = AMUpdate_gen(eye(3), 2.4/sqrt(3), 0.2, 1., 0.8, 25)
problem.adaptive_update = AMUpdate_gen(eye(3), 2.4/sqrt(3), 0.3, 1., 0.8, 25)

# set algorithm parameters
problem.alg_param.N = 1000 # nbr particels
problem.alg_param.R = 50000 # nbr iterations
problem.alg_param.burn_in = 0 # burn in
problem.alg_param.length_training_data = 5000
problem.alg_param.alg = "MCWM" # we should only! use the MCWM algorithm
problem.alg_param.compare_GP_and_PF = false
problem.alg_param.noisy_est = false
problem.alg_param.pred_method = "sample"
problem.alg_param.print_interval = 10000 # problem.alg_param.R#
problem.alg_param.beta_MH = 0.15 # "local_loglik_approx" # "max_loglik"

#problem.data.y = Array(readtable("y.csv"))[:,1]
#problem.data.y = Array(readtable("y_data_set_1.csv"))[:,1]
problem.data.y = Array(readtable("Ricker model/y_data_set_2.csv"))[:,1]

################################################################################
##                         training data                                      ##
################################################################################

problem_training = set_up_problem(ploton = false)

length_training_data = 2000
length_test_data = 2000
burn_in = 2000

problem_training.alg_param.N = 1000 # nbr particels
problem_training.alg_param.R = length_training_data + length_test_data + burn_in # nbr iterations
problem_training.alg_param.burn_in = burn_in # burn_in
problem_training.data.y = Array(readtable("Ricker model/y_data_set_2.csv"))[:,1] #Array(readtable("y.csv"))[:,1]
problem_training.alg_param.print_interval = 1000

# test starting at true parameters
#problem.model_param.theta_0 = problem.model_param.theta_true

# PMCMC
problem_training.alg_param.alg = "MCWM"

# use AM alg for adaptive updating
#problem.adaptive_update = AMUpdate(eye(3), 2.4/sqrt(3), 1., 0.7, 25)

#problem.adaptive_update = noAdaptation(2.4/sqrt(3)*eye(3))

# or, use AM gen alg for adaptive updating
#problem_training.adaptive_update = AMUpdate_gen(eye(3), 2.4/sqrt(3), 0.2, 1., 0.8, 25)
problem_training.adaptive_update = AMUpdate_gen(eye(3), 2.4/sqrt(3), 0.4, 1., 0.8, 25)

load_training_data = true

################################################################################
##                generate training data                                     ###
################################################################################

if !load_training_data

    # we can only load the training and test data

else

  #@load "gp_training_$(set_nbr_params)_par.jld"
  #@load "gp_training_$(set_nbr_params)_par.jld"

  @load "Ricker model/gp_training_and_test_data.jld"

end

################################################################################
###            plot training data                                            ###
################################################################################

#export_data(problem_training, res_training[1],"dagpMCMC_training_data"*jobname)

# plot chains for tranining
text_size = 15

PyPlot.figure(figsize=(10,20))

ax1 = PyPlot.subplot(311)
PyPlot.plot(res_training[1].Theta_est[1,:])
PyPlot.plot(problem.model_param.theta_true[1]*ones(size(res_training[1].Theta_est)[2]), "k")
PyPlot.ylabel(L"$\log r$",fontsize=text_size)
ax1[:axes][:get_xaxis]()[:set_ticks]([])

ax2 = PyPlot.subplot(312)
PyPlot.plot(res_training[1].Theta_est[2,:])
PyPlot.plot(problem.model_param.theta_true[2]*ones(size(res_training[1].Theta_est)[2]), "k")
PyPlot.ylabel(L"$\log \phi$",fontsize=text_size)
ax2[:axes][:get_xaxis]()[:set_ticks]([])

PyPlot.subplot(313)
PyPlot.plot(res_training[1].Theta_est[3,:])
PyPlot.plot(problem.model_param.theta_true[3]*ones(size(res_training[1].Theta_est)[2]), "k")
PyPlot.ylabel(L"$\log \sigma$",fontsize=text_size)
PyPlot.xlabel("Iteration",fontsize=text_size)

# plot chains after burn-in for tranining

PyPlot.figure(figsize=(10,20))

ax1 = PyPlot.subplot(311)
PyPlot.plot(res_training[1].Theta_est[1,burn_in:end])
PyPlot.plot(problem.model_param.theta_true[1]*ones(size(res_training[1].Theta_est[:,burn_in:end])[2]), "k")
PyPlot.ylabel(L"$\log r$",fontsize=text_size)
ax1[:axes][:get_xaxis]()[:set_ticks]([])

ax2 = PyPlot.subplot(312)
PyPlot.plot(res_training[1].Theta_est[2,burn_in:end])
PyPlot.plot(problem.model_param.theta_true[2]*ones(size(res_training[1].Theta_est[:,burn_in:end])[2]), "k")
PyPlot.ylabel(L"$\log \phi$",fontsize=text_size)
ax2[:axes][:get_xaxis]()[:set_ticks]([])

PyPlot.subplot(313)
PyPlot.plot(res_training[1].Theta_est[3,burn_in:end])
PyPlot.plot(problem.model_param.theta_true[3]*ones(size(res_training[1].Theta_est[:,burn_in:end])[2]), "k")
PyPlot.ylabel(L"$\log \sigma$",fontsize=text_size)
PyPlot.xlabel("Iteration",fontsize=text_size)

for i = 1:3
  PyPlot.figure()
  PyPlot.plot(res_training[1].Theta_est[i,:])
  PyPlot.plot(problem.model_param.theta_true[i]*ones(size(res_training[1].Theta_est)[2]), "k")
end

for i = 1:3
  PyPlot.figure()
  PyPlot.plot(data_training_star[i,:])
  PyPlot.plot(problem.model_param.theta_true[i]*ones(size(data_training_star,2)), "k")
end

for i = 1:3
  #PyPlot.figure()
  #PyPlot.plt[:hist](theta_training[i,:],100)
  PyPlot.figure()
  h = PyPlot.plt[:hist](data_training_star[i,:],100)
  PyPlot.plot((problem.model_param.theta_true[i], problem.model_param.theta_true[i]), (0, maximum(h[1])+5), "k");
end

PyPlot.figure()
PyPlot.plt[:hist](data_training_star[end,:],100)

PyPlot.figure()
PyPlot.plot(res_training[1].loglik_est)

show(mean(data_training_star[1:3,:],2))
show(std(data_training_star[1:3,:],2))

################################################################################
###          Fit GP model                                                     ##
################################################################################

# create gp object
gp = GPModel("est_method",zeros(6), zeros(4),
eye(problem.alg_param.length_training_data-20), zeros(problem.alg_param.length_training_data-20),zeros(2,problem.alg_param.length_training_data-20),
collect(1:10))

data_training = data_training_star
data_test = data_test_star

tic()

# fit GP model
if problem.alg_param.est_method == "ml"
  # fit GP model using ml

  perc_outlier = 0.1 # used when using PMCMC for trainig data 0.05
  tail_rm = "left"
  problem.alg_param.lasso = false

  ml_est(gp, data_training,"SE", problem.alg_param.lasso,perc_outlier,tail_rm)
else
  error("The two stage estimation method is not in use")
  #two_stage_est(gp, data_training)
end

time_fit_gp = toc()

################################################################################
##  Compare loglik predictions                                                                           ##
################################################################################

loglik_pf = data_test[end,:]

(loglik_mean,loglik_std,loglik_sample) = predict(data_test[1:end-1,:], gp, problem.alg_param.noisy_est)

# fit plot

# plot param vs loglik values
PyPlot.figure(figsize=(10,5))
PyPlot.plot(loglik_pf, "*b",alpha=0.3)
PyPlot.plot(loglik_sample, "*r",alpha=0.3)
PyPlot.ylabel(L"$\ell$",fontsize=text_size)
PyPlot.xlabel(L"Index",fontsize=text_size)

# plot param vs loglik values

for i = 1:3
  PyPlot.figure()
  PyPlot.plot(data_test[i,:], loglik_pf, "*b",alpha=0.3)
  PyPlot.plot(data_test[i,:], loglik_sample, "*r",alpha=0.3)
end

for i = 1:3
  PyPlot.figure()
  PyPlot.plot(data_test[i,:], loglik_pf, "*b",alpha=0.3)
  PyPlot.plot(data_test[i,:], loglik_mean, "*r",alpha=0.3)
end

# compute RMSE

RMSE_ml_mean = RMSE(loglik_pf, loglik_mean)
RMSE_ml_sample = RMSE(loglik_pf, loglik_sample)

################################################################################
##  Compute residuals                                                                          ##
################################################################################

residuals = loglik_pf - loglik_sample

# plot residuals
PyPlot.figure(figsize=(10,5))
PyPlot.plot(residuals)
PyPlot.ylabel("Residual",fontsize=text_size)
PyPlot.xlabel("Index",fontsize=text_size)

PyPlot.figure(figsize=(7,5))
PyPlot.plot(data_test[1,:], residuals, "*b")
PyPlot.ylabel("Residual",fontsize=text_size)
PyPlot.xlabel(L"$\log r$",fontsize=text_size)

PyPlot.figure(figsize=(7,5))
PyPlot.plot(data_test[2,:], residuals, "*b")
PyPlot.ylabel("Residual",fontsize=text_size)
PyPlot.xlabel(L"$\log \phi$",fontsize=text_size)

PyPlot.figure(figsize=(7,5))
PyPlot.plot(data_test[3,:], residuals, "*b")
PyPlot.ylabel("Residual",fontsize=text_size)
PyPlot.xlabel(L"$\log \sigma$",fontsize=text_size)

PyPlot.figure(figsize=(7,5))
PyPlot.plot(data_test[4,:], residuals, "*b")
PyPlot.ylabel("Residual",fontsize=text_size)
PyPlot.xlabel(L"$\ell$",fontsize=text_size)

PyPlot.figure(figsize=(7,5))
h1 = PyPlot.plt[:hist](residuals,100, normed=true)
PyPlot.xlabel("Residual",fontsize=text_size)
PyPlot.ylabel("Freq.",fontsize=text_size)

normplot(residuals)

################################################################################
##  Plot marginal functions pf as function of parameter value                                                                          ##
################################################################################

y = problem.data.y
theta_true = problem.model_param.theta_true
theta_known = problem.model_param.theta_known
N = 1000  #problem.alg_param.N
prior_parameters = problem.prior_dist.prior_parameters

include("pf.jl")

N = 1000

PyPlot.figure()
r_vec = prior_parameters[1,1]:0.01:prior_parameters[1,2]

loglik_r_non_fixed = zeros(2,length(r_vec))

for i = 1:length(r_vec)
  loglik_r_non_fixed[1,i] = pf(y, [r_vec[i]; theta_true[2:3]],theta_known,N,false)
  (loglik_mean,loglik_std,loglik_sample) = predict([r_vec[i]; theta_true[2:3]], gp, problem.alg_param.noisy_est)
  loglik_r_non_fixed[2,i]  = loglik_sample[1]
end

PyPlot.figure()
PyPlot.plot(r_vec, loglik_r_non_fixed[1,:], "b")
PyPlot.plot(r_vec, loglik_r_non_fixed[2,:], "r")
PyPlot.plot((theta_true[1], theta_true[1]), (minimum(loglik_r_non_fixed[find(!isnan, loglik_r_non_fixed)]), maximum(loglik_r_non_fixed[find(!isnan, loglik_r_non_fixed)])), "k")
PyPlot.plot((maximum(data_training[1,:]), maximum(data_training[1,:])), (minimum(loglik_r_non_fixed[find(!isnan, loglik_r_non_fixed)]), maximum(loglik_r_non_fixed[find(!isnan, loglik_r_non_fixed)])), "k--")
PyPlot.plot((minimum(data_training[1,:]), minimum(data_training[1,:])), (minimum(loglik_r_non_fixed[find(!isnan, loglik_r_non_fixed)]), maximum(loglik_r_non_fixed[find(!isnan, loglik_r_non_fixed)])), "k--")
PyPlot.xlabel(L"$\log r$")
PyPlot.ylabel(L"$\ell$")

# keeping psi non fixed

psi_vec = prior_parameters[2,1]:0.01:prior_parameters[2,2]

loglik_psi_non_fixed = zeros(2,length(psi_vec))

for i = 1:length(psi_vec)
  loglik_psi_non_fixed[1,i] = pf(y, [theta_true[1];psi_vec[i]; theta_true[3]],theta_known,N,false)
  (loglik_mean,loglik_std,loglik_sample) = predict( [theta_true[1];psi_vec[i]; theta_true[3]], gp, problem.alg_param.noisy_est)
  loglik_psi_non_fixed[2,i]  = loglik_sample[1]
end

PyPlot.figure()
PyPlot.plot(psi_vec, loglik_psi_non_fixed[1,:], "b")
PyPlot.plot(psi_vec, loglik_psi_non_fixed[2,:], "r")
PyPlot.plot((theta_true[2], theta_true[2]), (minimum(loglik_psi_non_fixed[find(!isnan, loglik_psi_non_fixed)]), maximum(loglik_r_non_fixed[find(!isnan, loglik_psi_non_fixed)])), "k")
PyPlot.plot((maximum(data_training[2,:]), maximum(data_training[2,:])), (minimum(loglik_psi_non_fixed[find(!isnan, loglik_psi_non_fixed)]), maximum(loglik_psi_non_fixed[find(!isnan, loglik_psi_non_fixed)])), "k--")
PyPlot.plot((minimum(data_training[2,:]), minimum(data_training[2,:])), (minimum(loglik_psi_non_fixed[find(!isnan, loglik_psi_non_fixed)]), maximum(loglik_psi_non_fixed[find(!isnan, loglik_psi_non_fixed)])), "k--")
PyPlot.xlabel(L"$\log \phi$")
PyPlot.ylabel(L"$\ell$")

# keeping psi non fixed

sigma_vec = prior_parameters[3,1]:0.01:prior_parameters[3,2]

loglik_sigma_non_fixed = zeros(2,length(sigma_vec))

for i = 1:length(sigma_vec)
  loglik_sigma_non_fixed[1,i] = pf(y,[theta_true[1:2];sigma_vec[i]],theta_known,N,false)
  (loglik_mean,loglik_std,loglik_sample) = predict([theta_true[1:2];sigma_vec[i]], gp, problem.alg_param.noisy_est)
  loglik_sigma_non_fixed[2,i]  = loglik_sample[1]
end

PyPlot.figure()
PyPlot.plot(sigma_vec, loglik_sigma_non_fixed[1,:], "b")
PyPlot.plot(sigma_vec, loglik_sigma_non_fixed[2,:], "r")
PyPlot.plot((theta_true[3], theta_true[3]), (minimum(loglik_sigma_non_fixed[find(!isnan, loglik_sigma_non_fixed)]), maximum(loglik_sigma_non_fixed[find(!isnan, loglik_sigma_non_fixed)])), "k")
PyPlot.plot((maximum(data_training[3,:]), maximum(data_training[3,:])), (minimum(loglik_sigma_non_fixed[find(!isnan, loglik_sigma_non_fixed)]), maximum(loglik_sigma_non_fixed[find(!isnan, loglik_sigma_non_fixed)])), "k--")
PyPlot.plot((minimum(data_training[3,:]), minimum(data_training[3,:])), (minimum(loglik_sigma_non_fixed[find(!isnan, loglik_sigma_non_fixed)]), maximum(loglik_sigma_non_fixed[find(!isnan, loglik_sigma_non_fixed)])), "k--")
PyPlot.xlabel(L"$\log \sigma$")
PyPlot.ylabel(L"$\ell$")

################################################################################
##  Selecting case problem                                                    ##
################################################################################

# Create features for classification models                                                                            ##

n = size(data_training,2)
n_burn_in = problem_training.alg_param.burn_in

dim = length(problem.model_param.theta_true)
data_gp_loglik_star_old = zeros(4,n)
data_gp_loglik_star_old[3,:] = data_training_star[dim+1,:]
data_gp_loglik_star_old[4,:] = data_training_old[dim+1,:]

# data_gp_loglik_star_old = [gp_star, gp_old, ll_star, ll_old]

std_pred_gp_star = zeros(n)

noisy_pred = problem.alg_param.noisy_est

for i = 1:n
  (loglik_est_star, var_pred_ml_star, prediction_sample_ml_star) = predict(data_training_star[1:dim,i],gp,noisy_pred)
  (loglik_est_old, var_pred_ml, prediction_sample_ml_old) = predict(data_training_old[1:dim,i],gp,noisy_pred)
  data_gp_loglik_star_old[1,i] = prediction_sample_ml_star[1]
  data_gp_loglik_star_old[2,i] = prediction_sample_ml_old[1]
  std_pred_gp_star[i] = sqrt(var_pred_ml_star[1])
end

nbr_GP_star_geq_GP_old = zero(Int64)
nbr_case_1 = zero(Int64)
nbr_case_2 = zero(Int64)

targets_case_1_and_3 = []
data_case_1_and_3 = []

targets_case_2_and_4 = []
data_case_2_and_4 = []

for i = 1:n
  if data_gp_loglik_star_old[1,i] > data_gp_loglik_star_old[2,i]
    nbr_GP_star_geq_GP_old += 1
    data_case_1_and_3 = vcat(data_case_1_and_3, [data_training[1:dim,i]; data_gp_loglik_star_old[1,i]/data_gp_loglik_star_old[2,i]; std_pred_gp_star[i]])
    if data_gp_loglik_star_old[3,i] > data_gp_loglik_star_old[4,i]
      append!(targets_case_1_and_3, 1)
      nbr_case_1 += 1
    else
      append!(targets_case_1_and_3, 0)
    end
  else
    data_case_2_and_4 = vcat(data_case_2_and_4, [data_training[1:dim,i]; data_gp_loglik_star_old[1,i]/data_gp_loglik_star_old[2,i]; std_pred_gp_star[i]])
    if data_gp_loglik_star_old[3,i] < data_gp_loglik_star_old[4,i]
      append!(targets_case_2_and_4, 1)
      nbr_case_2 += 1
    else
      append!(targets_case_2_and_4, 0)
    end
  end
end

# tansform features and set input data

# convert matricies to floats
data_case_1_and_3 = convert(Array{Float64,2},reshape(data_case_1_and_3, (dim+2, length(targets_case_1_and_3))))
data_case_2_and_4 = convert(Array{Float64,2},reshape(data_case_2_and_4, (dim+2, length(targets_case_2_and_4))))

targets_case_1_and_3 = convert(Array{Float64,1}, targets_case_1_and_3)
targets_case_2_and_4 = convert(Array{Float64,1}, targets_case_2_and_4)

# use untransformed features (biased coin model)

# is is not nessesary to construct input_data_case_1_and_3 and input_data_case_2_and_4 to run the biased coin model

input_data_case_1_and_3 = zeros(length(targets_case_1_and_3), dim+1)

input_data_case_1_and_3[:,1] =  data_case_1_and_3[1,:]
input_data_case_1_and_3[:,2] =  data_case_1_and_3[2,:]
input_data_case_1_and_3[:,3] =  data_case_1_and_3[3,:]
input_data_case_1_and_3[:,end] = targets_case_1_and_3

input_data_case_2_and_4 = zeros(length(targets_case_2_and_4), dim+1)

input_data_case_2_and_4[:,1] =  data_case_2_and_4[1,:]
input_data_case_2_and_4[:,2] =  data_case_2_and_4[2,:]
input_data_case_2_and_4[:,3] =  data_case_2_and_4[3,:]
input_data_case_2_and_4[:,end] = targets_case_2_and_4

# use tranformed data + standardization (logistic regression model)

# tansformed data for simple logistic model
mean_posterior = mean(theta_training,2)

input_data_case_1_and_3 = zeros(length(targets_case_1_and_3), dim+2)
input_data_case_1_and_3[:,1] = sqrt((mean_posterior[1] - data_case_1_and_3[1,:]).^2)
input_data_case_1_and_3[:,2] = sqrt((mean_posterior[2] - data_case_1_and_3[2,:]).^2)
input_data_case_1_and_3[:,3] = sqrt((mean_posterior[3] - data_case_1_and_3[3,:]).^2)
input_data_case_1_and_3[:,4] = sqrt(sum((repmat(mean_posterior', size(data_case_1_and_3,2))'-data_case_1_and_3[1:3,:]).^2,1))
input_data_case_1_and_3[:,end] = targets_case_1_and_3

input_data_case_2_and_4 = zeros(length(targets_case_2_and_4), dim+2)
input_data_case_2_and_4[:,1] = sqrt((mean_posterior[1] - data_case_2_and_4[1,:]).^2)
input_data_case_2_and_4[:,2] = sqrt((mean_posterior[2] - data_case_2_and_4[2,:]).^2)
input_data_case_2_and_4[:,3] = sqrt((mean_posterior[3] - data_case_2_and_4[3,:]).^2)
input_data_case_2_and_4[:,4] = data_case_1_and_3[4,:]
input_data_case_2_and_4[:,end] = targets_case_2_and_4

# tansformed data for logistic model with the ratio of gp ests are included

input_data_case_1_and_3 = zeros(length(targets_case_1_and_3), dim+3)
input_data_case_1_and_3[:,1] = sqrt((mean_posterior[1] - data_case_1_and_3[1,:]).^2)
input_data_case_1_and_3[:,2] = sqrt((mean_posterior[2] - data_case_1_and_3[2,:]).^2)
input_data_case_1_and_3[:,3] = sqrt((mean_posterior[3] - data_case_1_and_3[3,:]).^2)
input_data_case_1_and_3[:,4] = sqrt(sum((repmat(mean_posterior', size(data_case_1_and_3,2))'-data_case_1_and_3[1:3,:]).^2,1))
input_data_case_1_and_3[:,5] = data_case_1_and_3[4,:]
input_data_case_1_and_3[:,end] = targets_case_1_and_3

input_data_case_2_and_4 = zeros(length(targets_case_2_and_4), dim+3)
input_data_case_2_and_4[:,1] = sqrt((mean_posterior[1] - data_case_2_and_4[1,:]).^2)
input_data_case_2_and_4[:,2] = sqrt((mean_posterior[2] - data_case_2_and_4[2,:]).^2)
input_data_case_2_and_4[:,3] = sqrt((mean_posterior[3] - data_case_2_and_4[3,:]).^2)
input_data_case_2_and_4[:,4] = sqrt(sum((repmat(mean_posterior', size(data_case_2_and_4,2))'-data_case_2_and_4[1:3,:]).^2,1))
input_data_case_2_and_4[:,5] = data_case_2_and_4[4,:]
input_data_case_2_and_4[:,end] = targets_case_2_and_4

# features for the dt model

#standardization!(data_case_1_and_3)

input_data_case_1_and_3 = zeros(length(targets_case_1_and_3), dim+2)

input_data_case_1_and_3[:,1] =  data_case_1_and_3[1,:]
input_data_case_1_and_3[:,2] =  data_case_1_and_3[2,:]
input_data_case_1_and_3[:,3] =  data_case_1_and_3[3,:]
input_data_case_1_and_3[:,4] =  data_case_1_and_3[4,:]
input_data_case_1_and_3[:,end] = targets_case_1_and_3

#standardization!(data_case_2_and_4)

input_data_case_2_and_4 = zeros(length(targets_case_2_and_4), dim+2)

input_data_case_2_and_4[:,1] =  data_case_2_and_4[1,:]
input_data_case_2_and_4[:,2] =  data_case_2_and_4[2,:]
input_data_case_2_and_4[:,3] =  data_case_2_and_4[3,:]
input_data_case_2_and_4[:,4] =  data_case_2_and_4[4,:]
input_data_case_2_and_4[:,end] = targets_case_2_and_4

# plot features in input data

PyPlot.figure()
PyPlot.plot3D(input_data_case_1_and_3[find(x -> x==0, targets_case_1_and_3),1],
              input_data_case_1_and_3[find(x -> x==0, targets_case_1_and_3),2],
              input_data_case_1_and_3[find(x -> x==0, targets_case_1_and_3),3],
              "*r")
PyPlot.plot3D(input_data_case_1_and_3[find(x -> x==1, targets_case_1_and_3),1],
              input_data_case_1_and_3[find(x -> x==1, targets_case_1_and_3),2],
              input_data_case_1_and_3[find(x -> x==1, targets_case_1_and_3),3],
              "*g")
PyPlot.xlabel(L"dist \log r dim")
PyPlot.ylabel(L"dist \log \phi dim")
PyPlot.zlabel(L"dist  \sigma dim")

PyPlot.figure()
PyPlot.plot3D(input_data_case_1_and_3[find(x -> x==0, targets_case_1_and_3),2],
              input_data_case_1_and_3[find(x -> x==0, targets_case_1_and_3),3],
              input_data_case_1_and_3[find(x -> x==0, targets_case_1_and_3),4],
              "*r")
PyPlot.plot3D(input_data_case_1_and_3[find(x -> x==1, targets_case_1_and_3),2],
              input_data_case_1_and_3[find(x -> x==1, targets_case_1_and_3),3],
              input_data_case_1_and_3[find(x -> x==1, targets_case_1_and_3),4],
              "*g")
PyPlot.xlabel(L"dist \log \phi dim")
PyPlot.ylabel(L"dist  \sigma dim")
PyPlot.zlabel(L"joint dist")

# for the model with ratio
PyPlot.figure()
PyPlot.plot3D(input_data_case_1_and_3[find(x -> x==0, targets_case_1_and_3),3],
              input_data_case_1_and_3[find(x -> x==0, targets_case_1_and_3),4],
              input_data_case_1_and_3[find(x -> x==0, targets_case_1_and_3),5],
              "*r")
PyPlot.plot3D(input_data_case_1_and_3[find(x -> x==1, targets_case_1_and_3),3],
              input_data_case_1_and_3[find(x -> x==1, targets_case_1_and_3),4],
              input_data_case_1_and_3[find(x -> x==1, targets_case_1_and_3),5],
              "*g")
PyPlot.xlabel(L"dist  \sigma dim")
PyPlot.ylabel(L"joint dist")
PyPlot.zlabel(L"ratio gp ests")

PyPlot.figure()
PyPlot.plot3D(input_data_case_2_and_4[find(x -> x==0, targets_case_2_and_4),1],
              input_data_case_2_and_4[find(x -> x==0, targets_case_2_and_4),2],
              input_data_case_2_and_4[find(x -> x==0, targets_case_2_and_4),3],
              "*r")
PyPlot.plot3D(input_data_case_2_and_4[find(x -> x==1, targets_case_2_and_4),1],
              input_data_case_2_and_4[find(x -> x==1, targets_case_2_and_4),2],
              input_data_case_2_and_4[find(x -> x==1, targets_case_2_and_4),3],
              "*g")
PyPlot.xlabel(L"dist \log r dim")
PyPlot.ylabel(L"dist \log \phi dim")
PyPlot.zlabel(L"dist  \sigma dim")

PyPlot.figure()
PyPlot.plot3D(input_data_case_2_and_4[find(x -> x==0, targets_case_2_and_4),2],
              input_data_case_2_and_4[find(x -> x==0, targets_case_2_and_4),3],
              input_data_case_2_and_4[find(x -> x==0, targets_case_2_and_4),4],
              "*r")
PyPlot.plot3D(input_data_case_2_and_4[find(x -> x==1, targets_case_2_and_4),2],
              input_data_case_2_and_4[find(x -> x==1, targets_case_2_and_4),3],
              input_data_case_2_and_4[find(x -> x==1, targets_case_2_and_4),4],
              "*g")
PyPlot.xlabel(L"dist \log \phi dim")
PyPlot.ylabel(L"dist  \sigma dim")
PyPlot.zlabel(L"joint dist")

# for the model with ratio
PyPlot.figure()
PyPlot.plot3D(input_data_case_2_and_4[find(x -> x==0, targets_case_2_and_4),3],
              input_data_case_2_and_4[find(x -> x==0, targets_case_2_and_4),4],
              input_data_case_2_and_4[find(x -> x==0, targets_case_2_and_4),5],
              "*r")
PyPlot.plot3D(input_data_case_2_and_4[find(x -> x==1, targets_case_2_and_4),3],
              input_data_case_2_and_4[find(x -> x==1, targets_case_2_and_4),4],
              input_data_case_2_and_4[find(x -> x==1, targets_case_2_and_4),5],
              "*g")
PyPlot.xlabel(L"dist  \sigma dim")
PyPlot.ylabel(L"joint dist")
PyPlot.zlabel(L"ratio gp ests")

# use standadized features (features for decision tree model)

# 5 features model

standardization!(data_case_1_and_3)

input_data_case_1_and_3 = zeros(length(targets_case_1_and_3), dim+3)

input_data_case_1_and_3[:,1] =  data_case_1_and_3[1,:]
input_data_case_1_and_3[:,2] =  data_case_1_and_3[2,:]
input_data_case_1_and_3[:,3] =  data_case_1_and_3[3,:]
input_data_case_1_and_3[:,4] =  data_case_1_and_3[4,:]
input_data_case_1_and_3[:,5] =  data_case_1_and_3[5,:]
input_data_case_1_and_3[:,end] = targets_case_1_and_3

standardization!(data_case_2_and_4)

input_data_case_2_and_4 = zeros(length(targets_case_2_and_4), dim+3)

input_data_case_2_and_4[:,1] =  data_case_2_and_4[1,:]
input_data_case_2_and_4[:,2] =  data_case_2_and_4[2,:]
input_data_case_2_and_4[:,3] =  data_case_2_and_4[3,:]
input_data_case_2_and_4[:,4] =  data_case_2_and_4[4,:]
input_data_case_2_and_4[:,5] =  data_case_2_and_4[5,:]
input_data_case_2_and_4[:,end] = targets_case_2_and_4

# 4 features model

standardization!(data_case_1_and_3)

input_data_case_1_and_3 = zeros(length(targets_case_1_and_3), dim+2)

input_data_case_1_and_3[:,1] =  data_case_1_and_3[1,:]
input_data_case_1_and_3[:,2] =  data_case_1_and_3[2,:]
input_data_case_1_and_3[:,3] =  data_case_1_and_3[3,:]
input_data_case_1_and_3[:,4] =  data_case_1_and_3[4,:]
input_data_case_1_and_3[:,end] = targets_case_1_and_3

standardization!(data_case_2_and_4)

input_data_case_2_and_4 = zeros(length(targets_case_2_and_4), dim+2)

input_data_case_2_and_4[:,1] =  data_case_2_and_4[1,:]
input_data_case_2_and_4[:,2] =  data_case_2_and_4[2,:]
input_data_case_2_and_4[:,3] =  data_case_2_and_4[3,:]
input_data_case_2_and_4[:,4] =  data_case_2_and_4[4,:]
input_data_case_2_and_4[:,end] = targets_case_2_and_4

# plot features in input data

PyPlot.figure()
PyPlot.plot3D(input_data_case_1_and_3[find(x -> x==0, targets_case_1_and_3),1],
              input_data_case_1_and_3[find(x -> x==0, targets_case_1_and_3),2],
              input_data_case_1_and_3[find(x -> x==0, targets_case_1_and_3),3],
              "*r")
PyPlot.plot3D(input_data_case_1_and_3[find(x -> x==1, targets_case_1_and_3),1],
              input_data_case_1_and_3[find(x -> x==1, targets_case_1_and_3),2],
              input_data_case_1_and_3[find(x -> x==1, targets_case_1_and_3),3],
              "*g")
PyPlot.xlabel(L"\log r")
PyPlot.ylabel(L"\log \phi")
PyPlot.zlabel(L"\sigma")

PyPlot.figure()
PyPlot.plot3D(input_data_case_1_and_3[find(x -> x==0, targets_case_1_and_3),2],
              input_data_case_1_and_3[find(x -> x==0, targets_case_1_and_3),3],
              input_data_case_1_and_3[find(x -> x==0, targets_case_1_and_3),4],
              "*r")
PyPlot.plot3D(input_data_case_1_and_3[find(x -> x==1, targets_case_1_and_3),2],
              input_data_case_1_and_3[find(x -> x==1, targets_case_1_and_3),3],
              input_data_case_1_and_3[find(x -> x==1, targets_case_1_and_3),4],
              "*g")
PyPlot.xlabel(L"\log \phi")
PyPlot.ylabel(L"\sigma")
PyPlot.zlabel(L"ratio gp loglik est")

PyPlot.figure()
PyPlot.plot3D(input_data_case_1_and_3[find(x -> x==0, targets_case_1_and_3),2],
              input_data_case_1_and_3[find(x -> x==0, targets_case_1_and_3),3],
              input_data_case_1_and_3[find(x -> x==0, targets_case_1_and_3),5],
              "*r")
PyPlot.plot3D(input_data_case_1_and_3[find(x -> x==1, targets_case_1_and_3),2],
              input_data_case_1_and_3[find(x -> x==1, targets_case_1_and_3),3],
              input_data_case_1_and_3[find(x -> x==1, targets_case_1_and_3),5],
              "*g")
PyPlot.xlabel(L"\sigma")
PyPlot.ylabel(L"ratio gp loglik est")
PyPlot.zlabel(L"std gp est")

PyPlot.figure()
PyPlot.plot3D(input_data_case_2_and_4[find(x -> x==0, targets_case_2_and_4),1],
              input_data_case_2_and_4[find(x -> x==0, targets_case_2_and_4),2],
              input_data_case_2_and_4[find(x -> x==0, targets_case_2_and_4),3],
              "*r")
PyPlot.plot3D(input_data_case_2_and_4[find(x -> x==1, targets_case_2_and_4),1],
              input_data_case_2_and_4[find(x -> x==1, targets_case_2_and_4),2],
              input_data_case_2_and_4[find(x -> x==1, targets_case_2_and_4),3],
              "*g")
PyPlot.xlabel(L"\log r")
PyPlot.ylabel(L"\log \phi")
PyPlot.zlabel(L"\sigma")

PyPlot.figure()
PyPlot.plot3D(input_data_case_2_and_4[find(x -> x==0, targets_case_2_and_4),2],
              input_data_case_2_and_4[find(x -> x==0, targets_case_2_and_4),3],
              input_data_case_2_and_4[find(x -> x==0, targets_case_2_and_4),4],
              "*r")
PyPlot.plot3D(input_data_case_2_and_4[find(x -> x==1, targets_case_2_and_4),2],
              input_data_case_2_and_4[find(x -> x==1, targets_case_2_and_4),3],
              input_data_case_2_and_4[find(x -> x==1, targets_case_2_and_4),4],
              "*g")
PyPlot.xlabel(L"\log \phi")
PyPlot.ylabel(L"\sigma")
PyPlot.zlabel(L"ratio gp loglik est")

PyPlot.figure()
PyPlot.plot3D(input_data_case_2_and_4[find(x -> x==0, targets_case_2_and_4),2],
              input_data_case_2_and_4[find(x -> x==0, targets_case_2_and_4),3],
              input_data_case_2_and_4[find(x -> x==0, targets_case_2_and_4),5],
              "*r")
PyPlot.plot3D(input_data_case_2_and_4[find(x -> x==1, targets_case_2_and_4),2],
              input_data_case_2_and_4[find(x -> x==1, targets_case_2_and_4),3],
              input_data_case_2_and_4[find(x -> x==1, targets_case_2_and_4),5],
              "*g")
PyPlot.xlabel(L"\sigma")
PyPlot.ylabel(L"ratio gp loglik est")
PyPlot.zlabel(L"std gp est")

# Plot features

PyPlot.figure()
PyPlot.plot(data_case_1_and_3[1,:], targets_case_1_and_3, "*")

PyPlot.figure()
PyPlot.plot(data_case_1_and_3[2,:], targets_case_1_and_3, "*")

PyPlot.figure()
PyPlot.plot(data_case_1_and_3[3,:], targets_case_1_and_3, "*")

PyPlot.figure()
PyPlot.plot(data_case_1_and_3[4,:], targets_case_1_and_3, "*")

PyPlot.figure()
PyPlot.plot(data_case_1_and_3[4,:], targets_case_1_and_3, "*")

PyPlot.figure()
PyPlot.plot(data_case_1_and_3[5,:], targets_case_1_and_3, "*")

PyPlot.figure()
PyPlot.plot3D(data_case_1_and_3[1,find(x -> x==0, targets_case_1_and_3)],
              data_case_1_and_3[2,find(x -> x==0, targets_case_1_and_3)],
              data_case_1_and_3[3,find(x -> x==0, targets_case_1_and_3)],
              "*r")
PyPlot.plot3D(data_case_1_and_3[1,find(x -> x==1, targets_case_1_and_3)],
              data_case_1_and_3[2,find(x -> x==1, targets_case_1_and_3)],
              data_case_1_and_3[3,find(x -> x==1, targets_case_1_and_3)],
              "*g")
PyPlot.xlabel(L"\log r")
PyPlot.ylabel(L"\log \phi")
PyPlot.zlabel(L"log \sigma")

PyPlot.figure()
PyPlot.plot3D(data_case_1_and_3[1,find(x -> x==0, targets_case_1_and_3)],
              data_case_1_and_3[2,find(x -> x==0, targets_case_1_and_3)],
              data_case_1_and_3[4,find(x -> x==0, targets_case_1_and_3)],
              "*r")
PyPlot.plot3D(data_case_1_and_3[1,find(x -> x==1, targets_case_1_and_3)],
              data_case_1_and_3[2,find(x -> x==1, targets_case_1_and_3)],
              data_case_1_and_3[4,find(x -> x==1, targets_case_1_and_3)],
              "*g")
PyPlot.xlabel(L"\log r")
PyPlot.ylabel(L"\log \phi")
PyPlot.zlabel(L"ratio")

PyPlot.figure()
PyPlot.plot3D(data_case_1_and_3[3,find(x -> x==0, targets_case_1_and_3)],
              data_case_1_and_3[4,find(x -> x==0, targets_case_1_and_3)],
              data_case_1_and_3[5,find(x -> x==0, targets_case_1_and_3)],
              "*r")
PyPlot.plot3D(data_case_1_and_3[3,find(x -> x==1, targets_case_1_and_3)],
              data_case_1_and_3[4,find(x -> x==1, targets_case_1_and_3)],
              data_case_1_and_3[5,find(x -> x==1, targets_case_1_and_3)],
              "*g")
PyPlot.xlabel(L"\log \phi")
PyPlot.ylabel(L"ratio")
PyPlot.zlabel(L"std gp pred")

PyPlot.figure()
PyPlot.plot(data_case_2_and_4[1,:], targets_case_2_and_4, "*")

PyPlot.figure()
PyPlot.plot(data_case_2_and_4[2,:], targets_case_2_and_4, "*")

PyPlot.figure()
PyPlot.plot(data_case_2_and_4[3,:], targets_case_2_and_4, "*")

PyPlot.figure()
PyPlot.plot(data_case_2_and_4[4,:], targets_case_2_and_4, "*")

PyPlot.figure()
PyPlot.plot(data_case_2_and_4[5,:], targets_case_2_and_4, "*")

PyPlot.figure()
PyPlot.plot3D(data_case_2_and_4[1,find(x -> x==0, targets_case_2_and_4)],
              data_case_2_and_4[2,find(x -> x==0, targets_case_2_and_4)],
              data_case_2_and_4[3,find(x -> x==0, targets_case_2_and_4)],
              "*r")
PyPlot.plot3D(data_case_2_and_4[1,find(x -> x==1, targets_case_2_and_4)],
              data_case_2_and_4[2,find(x -> x==1, targets_case_2_and_4)],
              data_case_2_and_4[3,find(x -> x==1, targets_case_2_and_4)],
              "*g")
PyPlot.xlabel(L"\log r")
PyPlot.ylabel(L"\log \phi")
PyPlot.zlabel(L"log \sigma")

PyPlot.figure()
PyPlot.plot3D(data_case_1_and_3[1,find(x -> x==0, targets_case_1_and_3)],
              data_case_1_and_3[2,find(x -> x==0, targets_case_1_and_3)],
              data_case_2_and_4[4,find(x -> x==0, targets_case_1_and_3)],
              "*r")
PyPlot.plot3D(data_case_1_and_3[1,find(x -> x==1, targets_case_1_and_3)],
              data_case_1_and_3[2,find(x -> x==1, targets_case_1_and_3)],
              data_case_2_and_4[4,find(x -> x==1, targets_case_1_and_3)],
              "*g")
PyPlot.xlabel(L"\log r")
PyPlot.ylabel(L"\log \phi")
PyPlot.zlabel(L"ratio")

PyPlot.figure()
PyPlot.plot3D(data_case_1_and_3[3,find(x -> x==0, targets_case_1_and_3)],
              data_case_1_and_3[4,find(x -> x==0, targets_case_1_and_3)],
              data_case_2_and_4[5,find(x -> x==0, targets_case_1_and_3)],
              "*r")
PyPlot.plot3D(data_case_1_and_3[3,find(x -> x==1, targets_case_1_and_3)],
              data_case_1_and_3[4,find(x -> x==1, targets_case_1_and_3)],
              data_case_2_and_4[5,find(x -> x==1, targets_case_1_and_3)],
              "*g")
PyPlot.xlabel(L"\log \phi")
PyPlot.ylabel(L"ratio")
PyPlot.zlabel(L"std gp pred")

################################################################################
##  Biased coin model                                                                          ##
################################################################################

nbr_GP_star_led_GP_old = n-nbr_GP_star_geq_GP_old

prob_case_1 = nbr_case_1/nbr_GP_star_geq_GP_old
prob_case_2 = (nbr_case_2)/nbr_GP_star_led_GP_old
prob_case_3 = 1-prob_case_1
prob_case_4 = (nbr_GP_star_led_GP_old-nbr_case_2)/nbr_GP_star_led_GP_old
prob_cases = [prob_case_1;prob_case_2;prob_case_3;prob_case_4]

# test biased coin model
n = size(data_test,2)

nbr_case_1_selceted = 0
nbr_case_2_selceted = 0
nbr_case_3_selceted = 0
nbr_case_4_selceted = 0

nbr_case_1_correct = 0
nbr_case_2_correct = 0
nbr_case_3_correct = 0
nbr_case_4_correct = 0

include("pf.jl")
N = 1000
y = problem_training.data.y

case_1_theta_val = []
case_1_assumption_holds = []

case_2_theta_val = []
case_2_assumption_holds = []

case_3_theta_val = []
case_3_assumption_holds = []

case_4_theta_val = []
case_4_assumption_holds = []

for i = 1:n

  theta_new = data_test_star[1:dim,i]
  theta_old = data_test_old[1:dim,i]

  (loglik_est_star, var_pred_ml, prediction_sample_ml_star) = predict(theta_new,gp,noisy_pred)
  (loglik_est_old, var_pred_ml, prediction_sample_ml_old) = predict(theta_old,gp,noisy_pred)

  loglik_pf_old = pf(y, theta_old,problem_training.model_param.theta_known,N,false)
  loglik_pf_new = pf(y, theta_new,problem_training.model_param.theta_known,N,false)

  if prediction_sample_ml_star[1] > prediction_sample_ml_old[1]
    go_to_case_1 = rand() < prob_case_1
    if go_to_case_1

      case_1_theta_val = vcat(case_1_theta_val, theta_new)

      nbr_case_1_selceted = nbr_case_1_selceted+1

      if loglik_pf_new > loglik_pf_old
        nbr_case_1_correct = nbr_case_1_correct+1
        append!(case_1_assumption_holds, 1)

      else
        append!(case_1_assumption_holds, 0)
      end

    else
      nbr_case_3_selceted = nbr_case_3_selceted+1
      case_3_theta_val = vcat(case_3_theta_val, theta_new)

      if loglik_pf_new < loglik_pf_old
        nbr_case_3_correct = nbr_case_3_correct+1
        append!(case_3_assumption_holds, 1)
      else
        append!(case_3_assumption_holds, 0)
      end

    end

  else

    go_to_case_2 = rand() < prob_case_2

    if go_to_case_2

      nbr_case_2_selceted = nbr_case_2_selceted+1
      case_2_theta_val = vcat(case_2_theta_val, theta_new)

      if loglik_pf_new < loglik_pf_old
        nbr_case_2_correct = nbr_case_2_correct+1
        append!(case_2_assumption_holds, 1)
      else
        append!(case_2_assumption_holds, 0)
      end

    else
      nbr_case_4_selceted = nbr_case_4_selceted+1
      case_4_theta_val = vcat(case_4_theta_val, theta_new)

      if loglik_pf_new > loglik_pf_old
        nbr_case_4_correct = nbr_case_4_correct+1
        append!(case_4_assumption_holds, 1)
      else
        append!(case_4_assumption_holds, 0)
      end

    end
  end
end

case_1_theta_val = reshape(case_1_theta_val, (dim, length(case_1_assumption_holds)))
case_2_theta_val = reshape(case_2_theta_val, (dim, length(case_2_assumption_holds)))
case_3_theta_val = reshape(case_3_theta_val, (dim, length(case_3_assumption_holds)))
case_4_theta_val = reshape(case_4_theta_val, (dim, length(case_4_assumption_holds)))

nbr_case_selceted = [nbr_case_1_selceted nbr_case_2_selceted nbr_case_3_selceted nbr_case_4_selceted]

nbr_case_correct = [nbr_case_1_correct nbr_case_2_correct nbr_case_3_correct nbr_case_4_correct]

prob_correct_given_selected = nbr_case_correct./nbr_case_selceted

PyPlot.figure()
PyPlot.plot3D(case_1_theta_val[1,find(x -> x==0, case_1_assumption_holds)],
              case_1_theta_val[2,find(x -> x==0, case_1_assumption_holds)],
              case_1_theta_val[3,find(x -> x==0, case_1_assumption_holds)],
              "*r")
PyPlot.plot3D(case_1_theta_val[1,find(x -> x==1, case_1_assumption_holds)],
              case_1_theta_val[2,find(x -> x==1, case_1_assumption_holds)],
              case_1_theta_val[3,find(x -> x==1, case_1_assumption_holds)],
              "*g")
PyPlot.xlabel(L"\log r")
PyPlot.ylabel(L"\log \phi")
PyPlot.zlabel(L"log \sigma")

PyPlot.figure()
PyPlot.plot3D(case_2_theta_val[1,find(x -> x==0, case_2_assumption_holds)],
              case_2_theta_val[2,find(x -> x==0, case_2_assumption_holds)],
              case_2_theta_val[3,find(x -> x==0, case_2_assumption_holds)],
              "*r")
PyPlot.plot3D(case_2_theta_val[1,find(x -> x==1, case_2_assumption_holds)],
              case_2_theta_val[2,find(x -> x==1, case_2_assumption_holds)],
              case_2_theta_val[3,find(x -> x==1, case_2_assumption_holds)],
              "*g")
PyPlot.xlabel(L"\log r")
PyPlot.ylabel(L"\log \phi")
PyPlot.zlabel(L"log \sigma")

PyPlot.figure()
PyPlot.plot3D(case_3_theta_val[1,find(x -> x==0, case_3_assumption_holds)],
              case_3_theta_val[2,find(x -> x==0, case_3_assumption_holds)],
              case_3_theta_val[3,find(x -> x==0, case_3_assumption_holds)],
              "*r")
PyPlot.plot3D(case_3_theta_val[1,find(x -> x==1, case_3_assumption_holds)],
              case_3_theta_val[2,find(x -> x==1, case_3_assumption_holds)],
              case_3_theta_val[3,find(x -> x==1, case_3_assumption_holds)],
              "*g")
PyPlot.xlabel(L"\log r")
PyPlot.ylabel(L"\log \phi")
PyPlot.zlabel(L"log \sigma")

PyPlot.figure()
PyPlot.plot3D(case_4_theta_val[1,find(x -> x==0, case_4_assumption_holds)],
              case_4_theta_val[2,find(x -> x==0, case_4_assumption_holds)],
              case_4_theta_val[3,find(x -> x==0, case_4_assumption_holds)],
              "*r")
PyPlot.plot3D(case_4_theta_val[1,find(x -> x==1, case_4_assumption_holds)],
              case_4_theta_val[2,find(x -> x==1, case_4_assumption_holds)],
              case_4_theta_val[3,find(x -> x==1, case_4_assumption_holds)],
              "*g")
PyPlot.xlabel(L"\log r")
PyPlot.ylabel(L"\log \phi")
PyPlot.zlabel(L"log \sigma")

################################################################################
##  Logistic regression model                                                                             ##
################################################################################

using GLM

# fit model for cases 2 and 4

input_data_case_1_and_3 = DataFrame(input_data_case_1_and_3)

log_reg_model_case_1_and_3 = glm(@formula(x5 ~ x1 + x2 + x3 + x4), input_data_case_1_and_3, Binomial(), LogitLink())

# extened model
log_reg_model_case_1_and_3 = glm(@formula(x6 ~ x1 + x2 + x3 + x4 + x5), input_data_case_1_and_3, Binomial(), LogitLink())

# fit model for cases 2 and 4

input_data_case_2_and_4 = DataFrame(input_data_case_2_and_4)

log_reg_model_case_2_and_4 = glm(@formula(x5 ~ x1 + x2 + x3 + x4), input_data_case_2_and_4, Binomial(), LogitLink())

# extened model
log_reg_model_case_2_and_4 = glm(@formula(x6 ~ x1 + x2 + x3 + x4 + x5), input_data_case_2_and_4, Binomial(), LogitLink())

# test on training data

pred_training_data_case_1_and_3 = GLM.predict(log_reg_model_case_1_and_3) #, DataFrame(data_case_1_and_3'))

PyPlot.figure()
PyPlot.plot(data_case_1_and_3[1,:], log(pred_training_data_case_1_and_3./(1-pred_training_data_case_1_and_3)), "*")

PyPlot.figure()
PyPlot.plot(data_case_1_and_3[2,:], log(pred_training_data_case_1_and_3./(1-pred_training_data_case_1_and_3)), "*")

PyPlot.figure()
PyPlot.plot(data_case_1_and_3[3,:], log(pred_training_data_case_1_and_3./(1-pred_training_data_case_1_and_3)), "*")

PyPlot.figure()
PyPlot.plot(data_case_1_and_3[1,:], pred_training_data_case_1_and_3, "*")

PyPlot.figure()
PyPlot.plot(data_case_1_and_3[2,:], pred_training_data_case_1_and_3, "*")

PyPlot.figure()
PyPlot.plot(data_case_1_and_3[3,:], pred_training_data_case_1_and_3, "*")

PyPlot.figure()
PyPlot.plot3D(data_case_1_and_3[1,:],data_case_1_and_3[2,:], pred_training_data_case_1_and_3, "*")

PyPlot.figure()
PyPlot.plot3D(data_case_1_and_3[1,:],data_case_1_and_3[3,:], pred_training_data_case_1_and_3, "*")

PyPlot.figure()
PyPlot.plot3D(data_case_1_and_3[2,:],data_case_1_and_3[3,:], pred_training_data_case_1_and_3, "*")

pred_training_data_case_2_and_4 = GLM.predict(log_reg_model_case_2_and_4)#, #DataFrame(data_case_2_and_4'))

PyPlot.figure()
PyPlot.plot(data_case_2_and_4[1,:], log(pred_training_data_case_2_and_4./(1-pred_training_data_case_2_and_4)), "*")

PyPlot.figure()
PyPlot.plot(data_case_2_and_4[2,:], log(pred_training_data_case_2_and_4./(1-pred_training_data_case_2_and_4)), "*")

PyPlot.figure()
PyPlot.plot(data_case_2_and_4[3,:], log(pred_training_data_case_2_and_4./(1-pred_training_data_case_2_and_4)), "*")

PyPlot.figure()
PyPlot.plot(data_case_2_and_4[1,:], pred_training_data_case_2_and_4, "*")

PyPlot.figure()
PyPlot.plot(data_case_2_and_4[2,:], pred_training_data_case_2_and_4, "*")

PyPlot.figure()
PyPlot.plot(data_case_2_and_4[3,:], pred_training_data_case_2_and_4, "*")

PyPlot.figure()
PyPlot.plot3D(data_case_2_and_4[1,:],data_case_2_and_4[2,:], pred_training_data_case_2_and_4, "*")

PyPlot.figure()
PyPlot.plot3D(data_case_2_and_4[1,:],data_case_2_and_4[3,:], pred_training_data_case_2_and_4, "*")

PyPlot.figure()
PyPlot.plot3D(data_case_2_and_4[2,:],data_case_2_and_4[3,:], pred_training_data_case_2_and_4, "*")

# test model on test data

doc"""
   predict(beta::Vector, theta::Vector)

Prediction for the logistic regression model at theta.
"""
function predict(beta::Vector, theta::Vector)

  p_hat =  exp(theta'*beta)./(1+exp(theta'*beta))
  return p_hat[1]

end

beta_case_1_and_3 = coef(log_reg_model_case_1_and_3)
beta_case_2_and_4 = coef(log_reg_model_case_2_and_4)

n = size(data_test,2)

nbr_case_1_selceted = 0
nbr_case_2_selceted = 0
nbr_case_3_selceted = 0
nbr_case_4_selceted = 0

nbr_case_1_correct = 0
nbr_case_2_correct = 0
nbr_case_3_correct = 0
nbr_case_4_correct = 0

include("pf.jl")
N = 1000
y = problem_traning.data.y

for i = 1:n

  (loglik_est_star, var_pred_gp_star, prediction_sample_ml_star) = predict(data_test[1:dim,i],gp,noisy_pred)
  (loglik_est_old, var_pred_ml, prediction_sample_ml_old) = predict(res_training[1].Theta_est[:,idx_test_old[i]],gp,noisy_pred)

  theta_old = res_training[1].Theta_est[:,idx_test_old[i]]
  loglik_pf_old = pf(y, theta_old,problem_traning.model_param.theta_known,N,false)

  theta_new = data_test[1:dim,i] #res_training[1].Theta_est[:,i+n_burn_in+size(data_training,2)]
  loglik_pf_new = pf(y, theta_new,problem_traning.model_param.theta_known,N,false)

  # tansformation of theta_new
  # transform theta_new to euclidian distance space
  #=
  # old model
  theta_new_log_reg_mod = zeros(dim+1)

  theta_new_log_reg_mod[1] = sqrt((mean_posterior[1] - theta_new[1]).^2)
  theta_new_log_reg_mod[2] = sqrt((mean_posterior[2] - theta_new[2]).^2)
  theta_new_log_reg_mod[3] = sqrt((mean_posterior[3] - theta_new[3]).^2)
  theta_new_log_reg_mod[4] = sqrt(sum((mean_posterior - theta_new).^2))

  =#

  # for extened model
  theta_new_log_reg_mod = zeros(dim+2)

  theta_new_log_reg_mod[1] = sqrt((mean_posterior[1] - theta_new[1]).^2)
  theta_new_log_reg_mod[2] = sqrt((mean_posterior[2] - theta_new[2]).^2)
  theta_new_log_reg_mod[3] = sqrt((mean_posterior[3] - theta_new[3]).^2)
  theta_new_log_reg_mod[4] = sqrt(sum((mean_posterior - theta_new).^2))
  theta_new_log_reg_mod[5] = prediction_sample_ml_star[1]/prediction_sample_ml_old[1]

  if prediction_sample_ml_star[1] > prediction_sample_ml_old[1]
    #prob_case_1 = predict(beta_case_1_and_3, [1;data_test[1:dim,i]])
    prob_case_1 = predict(beta_case_1_and_3, [1;theta_new_log_reg_mod])

    go_to_case_1 = rand(Bernoulli(prob_case_1)) == 1

    if go_to_case_1

      nbr_case_1_selceted = nbr_case_1_selceted+1

      if loglik_pf_new > loglik_pf_old
        nbr_case_1_correct = nbr_case_1_correct+1
      end

    else
      nbr_case_3_selceted = nbr_case_3_selceted+1

      if loglik_pf_new < loglik_pf_old
        nbr_case_3_correct = nbr_case_3_correct+1
      end

    end

  else

    #prob_case_2 = predict(beta_case_2_and_4, [1;data_test[1:dim,i]])
    prob_case_2 = predict(beta_case_2_and_4, [1;theta_new_log_reg_mod])
    go_to_case_2 = rand(Bernoulli(prob_case_2)) == 1

    if go_to_case_2

      nbr_case_2_selceted = nbr_case_2_selceted+1

      if loglik_pf_new < loglik_pf_old
        nbr_case_2_correct = nbr_case_2_correct+1
      end

    else
      nbr_case_4_selceted = nbr_case_4_selceted+1

      if loglik_pf_new > loglik_pf_old
        nbr_case_4_correct = nbr_case_4_correct+1
      end

    end
  end
end

nbr_case_selceted = [nbr_case_1_selceted nbr_case_2_selceted nbr_case_3_selceted nbr_case_4_selceted]

nbr_case_correct = [nbr_case_1_correct nbr_case_2_correct nbr_case_3_correct nbr_case_4_correct]

prob_correct_given_selected = nbr_case_correct./nbr_case_selceted

################################################################################
##  Classification tree                                                                              ##
################################################################################

# Classification tree

using DecisionTree

# tree based model for case 1 and 3

# 5 feature model
features_1_and_3 = convert(Array, input_data_case_1_and_3[:, 1:dim+2])

# 4 features model
features_1_and_3 = convert(Array, input_data_case_1_and_3[:, 1:dim+1])

labels_1_and_3 = convert(Array, input_data_case_1_and_3[:, end])

labels_1_and_3 = Array{String}(size(features_1_and_3,1))

for i = 1:length(labels_1_and_3)
  if input_data_case_1_and_3[i,end] == 0
    labels_1_and_3[i] = "case 3"
  else
    labels_1_and_3[i] = "case 1"
  end
end

model_1_and_3 = build_tree(labels_1_and_3, features_1_and_3)

model_1_and_3 = prune_tree(model_1_and_3, 0.9)

print_tree(model_1_and_3, 3)

apply_tree(model_1_and_3, features_1_and_3[1,:])

accuracy = nfoldCV_tree(labels_1_and_3, features_1_and_3, 0.9, 3)

# tree based model for case 2 and 4

# 5 features model
features_case_2_and_4 = convert(Array, input_data_case_2_and_4[:, 1:dim+2])

# 4 features model
features_case_2_and_4 = convert(Array, input_data_case_2_and_4[:, 1:dim+1])

labels_case_2_and_4 = Array{String}(size(features_case_2_and_4,1))

for i = 1:length(labels_case_2_and_4)
  if input_data_case_2_and_4[i,end] == 0
    labels_case_2_and_4[i] = "case 4"
  else
    labels_case_2_and_4[i] = "case 2"
  end
end

# train full-tree classifier
model_2_and_4 = build_tree(labels_case_2_and_4, features_case_2_and_4)

model_2_and_4 = prune_tree(model_2_and_4, 0.9)

print_tree(model_2_and_4, 3)

apply_tree(model_2_and_4, features_case_2_and_4[1,:])

accuracy = nfoldCV_tree(labels_case_2_and_4, features_case_2_and_4, 0.9, 3)

# test model on test data

n = size(data_test,2)

nbr_case_1_selceted = 0
nbr_case_2_selceted = 0
nbr_case_3_selceted = 0
nbr_case_4_selceted = 0

nbr_case_1_correct = 0
nbr_case_2_correct = 0
nbr_case_3_correct = 0
nbr_case_4_correct = 0

include("pf.jl")
N = 1000
y = problem_training.data.y

for i = 1:n

  theta_new = data_test_star[1:dim,i]
  theta_old = data_test_old[1:dim,i]

  (loglik_est_star, var_pred_ml, prediction_sample_ml_star) = predict(theta_new,gp,noisy_pred)
  (loglik_est_old, var_pred_ml, prediction_sample_ml_old) = predict(theta_old,gp,noisy_pred)

  loglik_pf_old = pf(y, theta_old,problem_training.model_param.theta_known,N,false)
  loglik_pf_new = pf(y, theta_new,problem_training.model_param.theta_known,N,false)

  # tansformation of theta_new
  # transform theta_new to euclidian distance space

  # 5 features model
  #theta_new_tree_model = zeros(dim+2)
  #theta_new_tree_model = [theta_new; prediction_sample_ml_star[1]/prediction_sample_ml_old[1]; sqrt(var_pred_gp_star[1])]

  # 4 features model
  theta_new_tree_model = zeros(dim+1)
  theta_new_tree_model = [theta_new; prediction_sample_ml_star[1]/prediction_sample_ml_old[1]]

  if prediction_sample_ml_star[1] > prediction_sample_ml_old[1]
    #prob_case_1 = predict(beta_case_1_and_3, [1;data_test[1:dim,i]])
    case = apply_tree(model_1_and_3, theta_new_tree_model)

    if case ==  "case 1"

      nbr_case_1_selceted = nbr_case_1_selceted+1

      if loglik_pf_new > loglik_pf_old
        nbr_case_1_correct = nbr_case_1_correct+1
      end

    else
      nbr_case_3_selceted = nbr_case_3_selceted+1

      if loglik_pf_new < loglik_pf_old
        nbr_case_3_correct = nbr_case_3_correct+1
      end

    end

  else

    case = apply_tree(model_2_and_4, theta_new_tree_model)

    if case == "case 2"

      nbr_case_2_selceted = nbr_case_2_selceted+1

      if loglik_pf_new < loglik_pf_old
        nbr_case_2_correct = nbr_case_2_correct+1
      end

    else
      nbr_case_4_selceted = nbr_case_4_selceted+1

      if loglik_pf_new > loglik_pf_old
        nbr_case_4_correct = nbr_case_4_correct+1
      end

    end
  end
end

nbr_case_selceted = [nbr_case_1_selceted nbr_case_2_selceted nbr_case_3_selceted nbr_case_4_selceted]

nbr_case_correct = [nbr_case_1_correct nbr_case_2_correct nbr_case_3_correct nbr_case_4_correct]

prob_correct_given_selected = nbr_case_correct./nbr_case_selceted

max_stars_repo_name: SamuelWiqvist/adamcmcpaper

max_issues_repo_name: SamuelWiqvist/adamcmcpaper

max_forks_repo_name: SamuelWiqvist/adamcmcpaper


text:
function perform_redistancing(D)
    """
        perform_redistancing - redistance a function

          D1 = perform_redistancing(D, options);

          Compute a signed distance function D1 that has the same 0 level set as D (2D-matrix).

          Copyright (c) 2007 Gabriel Peyre
    """

    n = size(D,1)

    # % horizontal
    P1 = D[1 : end - 1, :]
    P2 = D[2 : end, :]
    P = (P1.*P2) .<= 0
    d = abs(P1-P2)

    l = collect((i, j) for i in 1:size(d)[1] for j in 1:size(d)[2] if d[i, j] < eps())
    for i in 1:size(d)[1]
        for j in 1:size(d)[2]
            if (i, j) in l
                d[i, j] =1
            end
        end
    end

    v1 = abs(P1)./d
    v2 = abs(P2)./d
    Ah = ([P; zeros(1,n)] + [zeros(1,n); P]) .> 0
    Vh = max([v1; zeros(1, n)], [zeros(1, n); v2])
    # % vertical
    P1 = D[:, 1 : end - 1]
    P2 = D[:, 2 : end]
    P = (P1.*P2) .<= 0;
    d = abs(P1-P2)
    l = collect((i, j) for i in 1:size(d)[1] for j in 1:size(d)[2] if d[i, j] < eps())
    for i in 1:size(d)[1]
        for j in 1:size(d)[2]
            if (i, j) in l
                d[i, j] =1
            end
        end
    end
    v1 = abs(P1)./d
    v2 = abs(P2)./d
    Av = ([P zeros(n, 1)] + [zeros(n, 1) P]) .> 0
    Vv = max([v1 zeros(n, 1)], [zeros(n,1) v2])

    V = zeros(n, n)
    I = findall(Ah .> 0)
    V[I] = Vh[I]
    I = findall(Av .> 0)
    V[I] = max(V[I], Vv[I])

    I = findall(V .!= 0)
    x, y = ind2sub(size(D), I)
    start_points = [x[:]'; y[:]']
    start_points[1, :], start_points[2, :] = start_points[2, :], start_points[1, :]

    D1 = NtToolBox.perform_fast_marching(ones(n, n), start_points)
    D1 = D1.*n
    D1[D .< 0] = -D1[D .< 0]

    return D1
end

ext: jl


text:
push!(LOAD_PATH, "../src/");

using Documenter, SocialDynamics

makedocs()

max_stars_repo_name: josephbakarji/socio-hydro-abm

max_issues_repo_name: josephbakarji/socio-hydro-abm

max_forks_repo_name: josephbakarji/socio-hydro-abm


text:
export dcat

"""
`dcat(A,B)` diagonally concatenate matrices `A` and `B`.
Returns a new matrix of the form `[A 0; 0 B]` where the `0`s
are blocks of zeros of the approriate shape.

More generally, may be called with any positive number of matrix, vector,
or scalar (numeric) arguments.
"""
function dcat(A::Array{S,2}, B::Array{T,2}) where {S,T}
    ST = promote_type(S, T)

    ra, ca = size(A)
    rb, cb = size(B)

    Z1 = zeros(ST, ra, cb)
    Z2 = zeros(ST, rb, ca)

    return [A Z1; Z2 B]
end

dcat(A) where {S} = _matrixify(A)

function dcat(A, B, C...)
    AB = dcat(A, B)
    return dcat(AB, C...)
end

function _matrixify(v::Array{T,1}) where {T}
    n = length(v)
    return reshape(v, n, 1)
end

_matrixify(x::T) where {T<:Number} = _matrixify([x])

_matrixify(A::Array{T,2}) where {T} = A

function dcat(v::T, A::Matrix{S}) where {S,T}
    return dcat(_matrixify(v), A)
end

function dcat(A::Matrix{S}, v::T) where {S,T}
    return dcat(A, _matrixify(v))
end

function dcat(x::T, y::S) where {S,T}
    return dcat(_matrixify(x), _matrixify(y))
end


text:
using DrWatson
quickactivate(@__DIR__)
@show projectdir()
##
using CovidSurvey
using StatsBase
using StatsPlots
using DataFrames
using Underscores
using HTTP
using Dates
using RData
using CSV
## ===========================================================================
# load data sets

fname = "/home/and/data/google-mobility/Global_Mobility_Report.csv"
mobil = CSV.File(fname) |> DataFrame

dateformat = "dduuuyyyy"
url = "https://raw.githubusercontent.com/OxCGRT/covid-policy-tracker/master/data/timeseries/stringency_index.csv"
stringency = @_ CSV.File(HTTP.get(url).body) |>
            DataFrame |>
            DataFrames.stack(__, Not([:Column1, :country_code, :country_name]), value_name=:stringency, variable_name=:date)
stringency[!,:date] = Date.(stringency.date, dateformat)

##
const mobility_names = [
    "retail_and_recreation_percent_change_from_baseline",
    "grocery_and_pharmacy_percent_change_from_baseline",
    "workplaces_percent_change_from_baseline",
    "residential_percent_change_from_baseline"
]

const new_mobility_names = [
    "retail",
    "grocery",
    "workplaces",
    "residential"
]
## ============================================================================
# dataset Denmark
mobil_dk = @_ mobil |>
                filter(:country_region => ==("Denmark"), __) |>
                filter(:sub_region_1 => ismissing, __) |>
                select(__, :date, mobility_names .=> new_mobility_names)
mobil_dk[!,:nonresidential] = [ mean(r) for r in eachrow( select(mobil_dk, Not([:date,:residential])) ) ]
@assert issorted(mobil_dk.date)

string_dk = @_ stringency |>
                filter(:country_name => ==("Denmark"), __) |>
                select(__, [:date, :stringency])

fname_epi = Dict(
    :cases => normpath( homedir(), "data/Statens-Serum-Institut/dashboard/Regionalt_DB/08_bekraeftede_tilfaelde_pr_dag_pr_regions.csv" ),
    :hospit => normpath( homedir(), "data/Statens-Serum-Institut/dashboard/Regionalt_DB/06_nye_indlaeggelser_pr_region_pr_dag.csv" ),
    :deaths => normpath( homedir(), "data/Statens-Serum-Institut/dashboard/Regionalt_DB/07_antal_doede_pr_dag_pr_region.csv" )
)

epi = Dict(
    n => @_ CSV.File( fn ) |>
            DataFrame |>
            National.process_ssi |>
            DataFrames.select(__, :date, :country => n) |>
            sort(__, :date)
    for (n, fn) in pairs(fname_epi)
)

data = let
    data = @_ RData.load("/home/and/code/epidemia/data/EuropeCovid2.RData")["EuropeCovid2"]["data"] |>
                filter(:country => ==("Denmark"), __) |>
                select(__, Not([:cases, :deaths, :id])) |>
                leftjoin(__, mobil_dk, on=:date) |>
                leftjoin(__, string_dk, on=:date) |>
                leftjoin(__, epi[:cases], on=:date) |>
                leftjoin(__, epi[:hospit], on=:date) |>
                leftjoin(__, epi[:deaths], on=:date) |>
                sort(__, :date)

    for n in [keys(epi)..., new_mobility_names..., "nonresidential"]
        data[!,n] = coalesce.(data[:,n], 0)
    end
    data[!,:pop] .= 5_837_213
    disallowmissing!(data)
    data
end

plot(data.residential)
plot!(data.nonresidential)

CSV.write("/home/and/Datathon/data/denmark.csv", data)

# f = KLogistic(.05)
# histogram( f.(rand(Normal(0,0.2), 10_000)) ,xticks=:native)
# histogram(rand(Normal(0.02,0.05), 10_000))
#
# let
#     plot(data.date, data.schools_universities, lab="schools")
#     plot!(data.date, data.public_events, lab="events")
#     plot!(data.date, data.self_isolating_if_ill, lab="isolating")
#     plot!(data.date, data.social_distancing_encouraged, lab="distancing")
#     plot!(data.date, data.lockdown, lab="lockdown")
# end
#
# plot(data.date, data.stringency, lab="stringency")
#
# let
#     plot(data.date, data.residential, lab="residential")
#     plot!(data.date, data.nonresidential, lab="nonresidential")
# end

## ============================================================================
# dataset Europe

data = RData.load("/home/and/code/epidemia/data/EuropeCovid2.RData")["EuropeCovid2"]["data"]
countries = unique(data.country)

function country_data(data, mobil, stringency, country)
    data_c  = filter(:country => ==(country), data)
    mobil_c = @_ mobil |>
                filter(:country_region => ==(country), __) |>
                filter(:sub_region_1 => ismissing, __) |>
                select(__, :date, mobility_names .=> new_mobility_names)
    mobil_c[!,:nonresidential] = [ mean(r) for r in eachrow( select(mobil_c, Not([:date,:residential])) ) ]
    @assert issorted(mobil_c.date)

    string_c = @_ stringency |>
                filter(:country_name => ==(country), __) |>
                select(__, [:date, :stringency])
    data_c = @_ data_c |>
                leftjoin(__, mobil_c, on=:date) |>
                leftjoin(__, string_c, on=:date)

    for n in [new_mobility_names..., "nonresidential"]
        data_c[!,n] = coalesce.(data_c[:,n], 0)
    end
    #disallowmissing!(data_c, :stringency)
    return data_c
end

newdata = let
    newdata = []
    for c in countries
        push!(newdata, country_data(data, mobil, stringency, c) )
    end
    vcat(newdata...)
end

CSV.write("/home/and/Datathon/data/EuropeCovid2.csv", newdata)

## ============================================================================
# regional hospitalizations

ps = (
    model = "hospit",
    observ = "2021-02-01",#"2021-02-06",#"2021-03-25"
    preds = nothing,#"CT,MRES,MNRES",#"CF,CC,CR,CS"
    semipara=false,
    rwstep = 7,
    include= true,
    std = true,
)

data_params = (
      observationsend  = ps.observ
    , predictors        = ps.preds|> CovidSurvey.parse_predictors
    # , cases_start       = ps.cases
    # , hospitmodel       = Regional.HospitInit1(obs_stop="2020-07-01")
    , casemodel         = Regional.CaseInit1(obs_start="2020-06-01")
    , seromodel         = Regional.SeroInit2(delay=0, std=1.)
    , rwstep            = ps.rwstep
    , epidemicstart    = 20
    , numimpute        = 6
    , include_early_dynamic = ps.include
    , link              = KLogistic(5.)
    , invlink           = KLogit(5.)
    , predictive        = false
    , covariates_kwargs = Dict(
      :semiparametric => false,
      :fname          => [projectdir("data", "mean_contact_rates_region=$(r).csv" ) for r in 1:Regional.nregions],#[projectdir("data", "mean_contact_rates_region=$r.csv" ) for r in 1:Regional.nregions],
      :shift          => -1,
      :startdate      => "2020-11-10", # >= 1000 cases / day
      :datecol        => "dates",
      :aggregation    => ps.rwstep,
      :standartize    => ps.std,
      :mobility       => [projectdir("data/google/mobility_region=$(Regional.regions[i]).csv") for i in 1:Regional.nregions]
      )
    )
latent_Rts_z[m] ~ RandomWalk(num_rt_steps[m])

data = Regional.load_data(; data_params... )
turing_data = data.turing_data;

using DataFrames
hospit = data.hospit
dates  = hospit.date

df = DataFrame(
      date=Vector{Date}()
    , week=Vector{Int64}()
    , hosp=Vector{Union{Int64, Missing}}()
    , region=Vector{String}()
)

for (i,(s,r)) in enumerate(zip(data.startdates, Regional.regions))
    mask = dates .>= s
    # date = (s-Day(30-1)):Day(1):last(dates) |> collect
    date = s:Day(1):last(dates) |> collect
    # hosp = vcat(fill(Missing, 30), hospit[mask, r])
    hosp = hospit[mask, r]
    region = fill(r, length(date))
    week = Regional.stepindex(length(date), 7)
    d    = DataFrame((; date, week, hosp, region))
    df = vcat(df, d)
end

hospits = copy(data.hospit)
hospits[!,:week] = Regional.stepindex(size(hospits,1), 7)
df = DataFrames.stack(hospits, Not([:week, :date]); variable_name=:region, value_name=:hospitalizations)
CSV.write("/home/and/Datathon/data/hopitalizations_regional.csv", df)

fdir = "/home/and/tmp/figures/"

using RData
fname = projectdir("data/COVID19/DK.rds")
c19 = RData.load(fname)

fname = projectdir("data/COVID19/DK.csv")
c19 = CSV.read(fname, DataFrame)
plot(c19.date, c19.confirmed)

df = copy(data.hospit)
df[!,:denmark] = sum(Array(df[:,Regional.regions]), dims=2 ) |> vec
plot!(data.deaths.date, df.denmark)

##                  OLD
# ## ============================================================================
# #                    data set for Aarhus Datathon
# #  ============================================================================
# using HTTP
# url = "https://raw.githubusercontent.com/OxCGRT/covid-policy-tracker/master/data/timeseries/stringency_index.csv"
# string = CSV.File(HTTP.get(url).body) |> DataFrame
#
# string = DataFrames.stack(string, Not([:Column1, :country_code, :country_name]), value_name=:stringency, variable_name=:date)
# string = filter(:country_name => ==("Denmark"), string)
# dateformat = "dduuuyyyy"
# string[!,:date] = Date.(string.date, dateformat)
# string = select(string, [:date, :stringency])
#
# fname_cases  = normpath( homedir(), "data/Statens-Serum-Institut/dashboard/Regionalt_DB/08_bekraeftede_tilfaelde_pr_dag_pr_regions.csv" )
# fname_hospit = normpath( homedir(), "data/Statens-Serum-Institut/dashboard/Regionalt_DB/06_nye_indlaeggelser_pr_region_pr_dag.csv" )
# fname_deaths = normpath( homedir(), "data/Statens-Serum-Institut/dashboard/Regionalt_DB/07_antal_doede_pr_dag_pr_region.csv" )
#
# @_ data |> DataFrames.select(__, :date)
#
# using Underscores
#
# cases = @_ CSV.File( fname_cases ) |>
#     DataFrame |>
#     National.process_ssi |>
#     DataFrames.select(__, :date, :country => :cases)
#
# hospit = @_ CSV.File( fname_hospit ) |>
#     DataFrame |>
#     National.process_ssi |>
#     DataFrames.select(__, :date, :country => :hospit)
#
# deaths = @_ CSV.File( fname_deaths ) |>
#     DataFrame |>
#     National.process_ssi |>
#     DataFrames.select(__, :date, :country => :deaths)
#
# using RData
# data = RData.load("/home/and/code/epidemia/data/EuropeCovid2.RData")["EuropeCovid2"]["data"]
# data = filter(:country => ==("Denmark"), data)
# data = select(data, Not([:cases, :deaths, :id, :country]))
# data = leftjoin(data, mobil_dk, on=:date)
# data = leftjoin(data, string, on=:date)
# data = leftjoin(data, cases, on=:date)
# data = leftjoin(data, hospit, on=:date)
# data = leftjoin(data, deaths, on=:date)
#
# CSV.write("/home/and/Datathon/data/denmark.csv", data)

max_stars_repo_name: andreaskoher/CovidSurvey

max_issues_repo_name: andreaskoher/CovidSurvey

max_forks_repo_name: andreaskoher/CovidSurvey


text:
const MASKCOL = ARGB32(0.5)
const ZEROCOL = ARGB32(0.3)

"""
    Renderer

Abstract supertype for objects that convert a frame of the simulation into an `ARGB32`
image for display. Frames may be a single grid or a `NamedTuple` of multiple grids.
"""
abstract type Renderer end

imagesize(r::Renderer, e::Extent) = imagesize(r, init(e), tspan(e))
imagesize(r::Renderer, init::NamedTuple, tspan) = imagesize(r, first(init), tspan)
imagesize(::Renderer, init::AbstractArray{<:Any,1}, tspan) = (length(tspan), size(init)...)
imagesize(::Renderer, init::AbstractArray, tspan) = size(init)

# 1D : timespan y axis gets filled during the simulation
function _allocimage(r::Renderer, init, tspan)
    fill(ARGB32(0), imagesize(r, init, tspan)...)
end

# render!
# Converts grid/s to an image for viewing.
function render!(o::Output, data::AbstractSimData)
    render!(o, data, grids(data))
end
function render!(o::Output, data::AbstractSimData, grids)
    render!(imagebuffer(o), renderer(o), o, data, grids)
end

# Single-grid renderers ###############################################################

"""
    SingleGridRenderer <: Renderer

Abstract supertype for [`Renderer`](@ref)s that convert a single grid
into an image array.

The first grid will be displayed if a `SingleGridRenderer` is
used with a `NamedTuple` of grids.
"""
abstract type SingleGridRenderer <: Renderer end

function render!(
    imagebuffer, ig::SingleGridRenderer, o::Output, data::AbstractSimData,
    grids::NamedTuple;
    name=string(first(keys(grids))), time=currenttime(data)
)
    render!(imagebuffer, ig, o, data, first(grids); name=name, time=time)
end
function render!(
    imagebuffer, ig::SingleGridRenderer, o::Output, data::AbstractSimData,
    grids::NamedTuple{(DEFAULT_KEY,)};
    name=nothing, time=currenttime(data)
)
    render!(imagebuffer, ig, o, data, first(grids); name=name, time=time)
end
function render!(
    imagebuffer, ig::SingleGridRenderer, o::Output,
    data::AbstractSimData{S}, grid::AbstractArray{<:Any,2};
    name=nothing, time=currenttime(data), accessor=nothing,
    minval=minval(o), maxval=maxval(o),
) where S<:Tuple{Y,X} where {Y,X}
    for j in 1:X, i in 1:Y
        accessor = accessor isa Nothing ? DynamicGrids.accessor(ig) : accessor
        @inbounds grid_obj = grid[i, j]
        val = _access_grid_object(accessor, grid_obj)
        pixel = to_rgb(cell_to_pixel(ig, mask(o), minval, maxval, data, val, (i, j)))
        @inbounds imagebuffer[i, j] = pixel
    end
    _rendertext!(imagebuffer, textconfig(o), name, time)
    return imagebuffer
end
function render!(
    imagebuffer, ig::SingleGridRenderer, o::GraphicOutput,
    data::AbstractSimData{S}, grid::AbstractArray{<:Any,1};
    name=nothing, time=currenttime(data), accessor=nothing,
    minval=minval(o), maxval=maxval(o),
) where S<:Tuple{L} where {L}
    f = currentframe(data)
    for i in 1:L
        accessor = accessor isa Nothing ? DynamicGrids.accessor(ig) : accessor
        @inbounds grid_obj = grid[i]
        val = _access_grid_object(accessor, grid_obj)
        pixel = to_rgb(cell_to_pixel(ig, mask(o), minval, maxval, data, val, (i,)))
        # Image rows are simulation frames
        @inbounds imagebuffer[f, i] = pixel
    end
    _rendertext!(imagebuffer, textconfig(o), name, time)
    return imagebuffer
end

const IMAGE_RENDERER_KEYWORDS = """
- `scheme`: a ColorSchemes.jl colorscheme, [`ObjectScheme`](@ref) or object that defines
    `Base.get(obj, val)` and returns a `Color` or a value that can be converted to `Color`
    using `ARGB32(val)`.
- `zerocolor`: a `Col` to use when values are zero, or `nothing` to ignore.
- `maskcolor`: a `Color` to use when cells are masked, or `nothing` to ignore.
"""

"""
    Image <: SingleGridRenderer

    Image(f=identity; scheme=ObjectScheme(), zerocolor=nothing, maskcolor=nothing)

Converts output grids to a colorsheme.

# Arguments

- `f`: a function to convert value from the grid to `Real`
    oran `RGB`. `Real` will be scaled by minval/maxval and be colored by the `scheme`.
    `RGB` is used directly in the output. This is useful for grids of complex objects,
    but not necessary for numbers. The default is `identity`.

# Keywords

$IMAGE_RENDERER_KEYWORDS
"""
struct Image{A,S,Z,M} <: SingleGridRenderer
    accessor::A
    scheme::S
    zerocolor::Z
    maskcolor::M
end
Image(accesor::Union{Function,Int}, scheme, zerocolor=ZEROCOL) = Image(scheme, zerocolor, MASKCOL)
Image(scheme, zerocolor=ZEROCOL, maskcolor=MASKCOL) = Image(identity, scheme, zerocolor, MASKCOL)
Image(accessor::Union{Function,Int}; kw...) = Image(; accessor, kw...)
Image(; accessor::Union{Function,Int}=identity, scheme=ObjectScheme(), zerocolor=ZEROCOL, maskcolor=MASKCOL, kw...) =
    Image(accessor, scheme, zerocolor, maskcolor)

accessor(p::Image) = p.accessor
scheme(p::Image) = p.scheme
zerocolor(p::Image) = p.zerocolor
maskcolor(p::Image) = p.maskcolor

# Show colorscheme in Atom etc
Base.show(io::IO, m::MIME"image/svg+xml", p::Image) = show(io, m, scheme(p))

# cell_to_pixel!
# Convert a single sell to an ARGB pixel
@inline function cell_to_pixel(ig::Image, mask, minval, maxval, data::AbstractSimData, val, I)
    if !(maskcolor(ig) isa Nothing) && ismasked(mask, I...)
        to_rgb(maskcolor(ig))
    else
        normval = normalise(val, minval, maxval)
        if !(zerocolor(ig) isa Nothing) && normval == zero(typeof(normval))
            to_rgb(zerocolor(ig))
        elseif normval isa Number && isnan(normval)
            zerocolor(ig) isa Nothing ? to_rgb(scheme(ig), 0) : to_rgb(zerocolor(ig))
        else
            to_rgb(scheme(ig), normval)
        end
    end
end

# Multi-grid renderers ###############################################################

"""
    MultiGridRenderer <: Renderer

Abstract type for `Renderer`s that convert a frame containing multiple
grids into a single image.
"""
abstract type MultiGridRenderer <: Renderer end

"""
    Layout <: MultiGridRenderer

    Layout(layout::Array, renderer::Matrix)

Layout allows displaying multiple grids in a block layout, by specifying a
layout matrix and a list of [`Image`](@ref)s to be run for each.

# Arguments

- `layout`: A `Vector` or `Matrix` containing the keys or numbers of grids in the
    locations to display them. `nothing`, `missing` or `0` values will be skipped.
- `renderers`: `Vector/Matrix` of [`Image`](@ref), matching the `layout`.
    Can be `nothing` or any other value for grids not in layout.
"""
Base.@kwdef struct Layout{L<:Union{AbstractVector,AbstractMatrix},R} <: MultiGridRenderer
    layout::L
    renderers::R
    Layout(layouts::L, renderers::R) where {L,R} = begin
        renderers1 = map(_asrenderer, renderers)
        new{L,typeof(renderers1)}(layouts, renderers1)
    end
end

layout(l::Layout) = l.layout
renderers(l::Layout) = l.renderers

imagesize(l::Layout, init::NamedTuple, tspan) = imagesize(l, first(init), tspan)
function imagesize(l::Layout, init::AbstractArray, tspan)
    _imagesize(size(init), size(l.layout), tspan)
end
# 1D
_imagesize(gsize::NTuple{1}, lsize::NTuple{1}, tspan) = length(tspan), gsize .* lsize
# 2D
_imagesize(gsize::NTuple{2}, lsize::NTuple{2}, tspan) = gsize .* lsize
_imagesize(gsize::NTuple{2}, lsize::NTuple{1}, tspan) = gsize .* (first(lsize), 1)

_asrenderer(r::Renderer) = r
_asrenderer(x) = Image(x)

function render!(
    imagebuffer, l::Layout, o::Output, data::AbstractSimData, grids::NamedTuple
)
    npanes = length(layout(l))
    minv, maxv = minval(o), maxval(o)
    if !(minv isa Union{Number,Nothing})
        length(minv) == npanes || _wronglengtherror(minval, npanes, length(minv))
    end
    if !(maxv isa Union{Number,Nothing})
        length(maxv) == npanes || _wronglengtherror(maxval, npanes, length(maxv))
    end

    grid_ids = map(_grid_ids, layout(l))
    grid_accessors = map(_grid_accessor, layout(l))
    # Loop over the layout matrix
    for I in CartesianIndices(grid_ids)
        grid_id = grid_ids[I]
        # Accept symbol keys and numbers, skip missing/nothing/0
        (ismissing(grid_id) || grid_id === nothing || grid_id == 0)  && continue
        n = if grid_id isa Symbol
            found = findfirst(k -> k === grid_id, keys(grids))
            found === nothing && _grididnotinkeyserror(grid_id, grids)
            found
        else
            grid_id
        end
        Itup = length(Tuple(I)) == 1 ? (Tuple(I)..., 1) : Tuple(I)
        im_I = map(Itup, gridsize(data)) do l, gs
            (l - 1) * gs + 1:l * gs
        end
        lin = LinearIndices(grid_ids)[I]

        # Run image renderers for section
        render!(
            view(imagebuffer, im_I...), renderers(l)[lin], o, data, grids[n];
            name=string(keys(grids)[n]), time=nothing,
            minval=_get(minv, lin), maxval=_get(maxv, lin),
            accessor=grid_accessors[I],
        )
    end
    _rendertime!(imagebuffer, textconfig(o), currenttime(data))
    return imagebuffer
end

_get(::Nothing, I) = nothing
_get(vals::Union{AbstractArray,Tuple}, I) = vals[I]
_get(x, I) = x

# Get the id of each grid in the layout
_grid_ids(id::Symbol) = id
_grid_ids(id::Integer) = id
_grid_ids(::Nothing) = nothing
_grid_ids(::Missing) = missing
_grid_ids(id::Pair{Symbol}) = first(id)
_grid_ids(id) = throw(ArgumentError("Layout id $id is not a valid grid name. Use an `Int`, `Symbol`, `Pair{Symbol,<:Any}` or `nothing`"))

# Get the grid accessor function/value if there is one
_grid_accessor(id::Pair) = last(id)
_grid_accessor(id) = nothing

# Access a grid object with a function or index, or use it as-is
_access_grid_object(::Nothing, obj) = obj
_access_grid_object(f::Function, obj) = f(obj)
_access_grid_object(i::Int, obj) = obj[i]

@noinline _grididnotinkeyserror(grid_id, grids) =
    throw(ArgumentError("$grid_id is not in $(keys(grids))"))
@noinline _wronglengtherror(f, npanes, len) =
    throw(ArgumentError("Number of layout panes ($npanes) and length of $f ($len) must be the same"))

# Automatically choose an image renderer ############################################################

autorenderer(init; scheme=ObjectScheme(), zerocolor=ZEROCOL, maskcolor=MASKCOL, kw...) =
    _autorenderer(init, scheme, zerocolor, maskcolor; kw...)

# Single grid render
_autorenderer(init, scheme, zerocolor, maskcolor; kw...) = _asrenderer(scheme, zerocolor, maskcolor)
# Multi-grid renders
function _autorenderer(init::NamedTuple, scheme, zerocolor, maskcolor;
    # If no layout argument was passed in, generate the layout from the init grids
    layout=_autolayout(init),
    # If no renders argument was passed in, generate them from the layout
    renderers=_asrenderer.(layout, _iterable(scheme), _iterable(zerocolor), _iterable(maskcolor)),
    kw...
)
    Layout(layout, renderers)
end

# Maybe wrap in a tuple for broadcasting
_iterable(obj::AbstractArray) = obj
_iterable(obj::Tuple) = obj
_iterable(obj) = (obj,)

# _asrenderer
# Use a renderer as-is or convert a colorscheme to a renderer
_asrenderer(renderer::Renderer, zerocolor, maskcolor) = renderer
_asrenderer(scheme, zerocolor, maskcolor) = Image(scheme, zerocolor, maskcolor)
_asrenderer(layout, renderer, zerocolor, maskcolor) = _asrenderer(renderer, zerocolor, maskcolor)

# _autolayout
# Generate a layout from the provided grids and their contents
function _autolayout(init)
    keys = _autokeys(init)
    len = length(keys)
    rows = len ÷ 4 + 1
    cols = (len - 1) ÷ rows + 1
    layout = Array{Any}(fill(nothing, rows, cols))
    for i in eachindex(keys)
        layout[i] = keys[i]
    end
    return layout
end

# _autokeys
# Generate a list of keys for layout grids.
# These may simply be the NamedTuple key or Pairs
# of the key name and index value, if the grid contains
# objects like FieldVector or StaticArray
function _autokeys(init)
    foldl(pairs(init); init=()) do acc, (key, val)
        keys = if eltype(val) <: AbstractArray
            cellcontents = first(val)
            map(i -> key => i, 1:length(cellcontents))
        else
            (key,)
        end
        (acc..., keys...)
    end
end

# Colour conversion tools ############################################333

# Set a value to be between zero and one, before converting to Color.
# min and max of `nothing` are assumed to be 0 and 1.
normalise(x::X, minv, maxv) where X = max(min((x - minv) / (maxv - minv), oneunit(X)), zero(X))
normalise(x::X, minv, maxv::Nothing) where X = max((x - minv) / (oneunit(X) - minv), zero(X))
normalise(x::X, minv::Nothing, maxv) where X = min(x / maxv, oneunit(X), oneunit(X))
normalise(x, minv::Nothing, maxv::Nothing) = x
normalise(x; min=Nothing, max=Nothing) = normalise(x, min, max)

# Rescale a value between 0 and 1 to be between `min` and `max`.
# This can be used to shrink the range of a colorsheme that is displayed.
# min and max of `nothing` are assumed to be 0 and 1.
scale(x, min, max) = x * (max - min) + min
scale(x, ::Nothing, max) = x * max
scale(x, min, ::Nothing) = x * (oneunit(typeof(min)) - min) + min
scale(x, ::Nothing, ::Nothing) = x

to_rgb(vals::Tuple) = ARGB32(vals...)
to_rgb(val::Real) = ARGB32(RGB(val))
to_rgb(val::Color) = ARGB32(val)
to_rgb(val::ARGB32) = val
# Handle external colorshemes, such as from ColorSchemes.jl
to_rgb(scheme, val::Real) = to_rgb(get(scheme, val))


text:
using Test

requiredFields = ["byr", "iyr", "eyr", "hgt", "hcl", "ecl", "pid"]

mutable struct Passport
    byr:: Int
    iyr:: Int
    eyr:: Int
    hgt:: Int
    hcl:: Int
    ecl:: Int
    pid:: Int
    isValid:: Bool
    Passport() = new()
end

function addLineToPassport(line, passport)
    fieldes = split(line," ")
    for field in fieldes
        (key::String, value:: String) = split(field, ":")
        if key == "byr" && isValidByr(value) passport.byr = 1 end
        if key == "iyr" && isValidIyr(value) passport.iyr = 1 end
        if key == "eyr" && isValidEyr(value) passport.eyr = 1 end
        if key == "hgt" && isValidHgt(value) passport.hgt = 1 end
        if key == "hcl" && isValidHcl(value) passport.hcl = 1 end
        if key == "ecl" && isValidEcl(value) passport.ecl = 1 end
        if key == "pid" && isValidPid(value) passport.pid = 1 end
    end
    return passport
end

function isValidNumberRange(value:: String, min::Int, max::Int, size::Int)
    isCorrectSize = length(value) == size
    intValue = 0
    try
        intValue = parse(Int, value)
    catch
        return false
    end
    isValidRange = (intValue >= min) && (intValue <= max)
    return  isCorrectSize && isValidRange
end

function isValidByr(value)
    return isValidNumberRange(value, 1920, 2002, 4)
end

function isValidIyr(value)
    return isValidNumberRange(value, 2010, 2020, 4)
end

function isValidEyr(value)
    return isValidNumberRange(value, 2020, 2030, 4)
end

function isValidHgt(valueString)
    result = false
    if(length(valueString) < 3)
        return false
    else
        unit = SubString(valueString, length(valueString)-1)
        value::String = SubString(valueString, 1:length(valueString)-2)
        if unit == "cm" && isValidNumberRange(value,150,193,3)
            result = true
        end
        if unit == "in" && isValidNumberRange(value,59,76,2)
            result = true
        end
    end
    return result
end

function isValidHcl(value)
    rx = r"(#[a-z0-9]{6})"
    matchResult = match(rx, value)
    return matchResult != nothing && matchResult.match == value
end

function isValidEcl(value)
    rx = r"amb|blu|brn|gry|grn|hzl|oth"
    matchResult = match(rx, value)
    return matchResult != nothing && matchResult.match == value
end

function isValidPid(value)
    rx = r"([0-9]{9})"
    matchResult = match(rx, value)
    return matchResult != nothing && matchResult.match == value
end

function validatePassport(p::Passport)
    if p.byr + p.iyr + p.eyr + p.hgt + p.hcl + p.ecl + p.pid == 7
        p.isValid = true
    end
end

function collectPassports(fileName)
    passports = []
    passportCounter = 1
    currentPassport = Passport()
    open(fileName) do file
        for line in eachline(file)
            if length(line) == 0
                validatePassport(currentPassport)
                push!(passports, currentPassport)
                currentPassport = Passport()
                passportCounter += 1
            else
                addLineToPassport(line, currentPassport)
            end
        end
    end
    return passports
end

allPassports = collectPassports("data.txt")
numberOfValidPassports = 0
for passport in allPassports
    if passport.isValid
        global numberOfValidPassports += 1
    end
end
println(numberOfValidPassports)

# line = "hcl:1f7352 iyr:2014 hgt:70cm eyr:1983"
# passport = Passport()
# println(passport)
# pasport = addLineToPassport(line, passport)
# println(passport)

# passports = collectPassports("data_test.txt")
# for passport in passports
#     println(passport)
# end

function assertTrue(value::Bool)
    return value == true
end

function assertFalse(value::Bool)
    return value == false
end

# println("Byr")
# println(assertFalse(isValidByr("1900")))
# println(assertFalse(isValidByr("2003")))
# println(assertFalse(isValidByr("abc")))
# println(assertFalse(isValidByr("201")))
# println(assertTrue(isValidByr("1972")))

# println("\nIyr")
# println(assertFalse(isValidIyr("abc")))
# println(assertFalse(isValidIyr("1972")))
# println(assertTrue(isValidIyr("2011")))
# println(assertFalse(isValidIyr("2040")))

# println("\nEyr")
# println(assertFalse(isValidEyr("abc")))
# println(assertFalse(isValidEyr("1972")))
# println(assertTrue(isValidEyr("2021")))
# println(assertFalse(isValidEyr("2040")))

# println("\nHgt")
# println(assertFalse(isValidHgt("123cm")))
# println(assertTrue(isValidHgt("151cm")))
# println(assertFalse(isValidHgt("151in")))
# println(assertTrue(isValidHgt("59in")))
# println(assertFalse(isValidHgt("77in")))

# println("\nHcl")
# println(assertFalse(isValidHcl("#1234a56.")))
# println(assertFalse(isValidHcl("#1234a51234a5")))
# println(assertTrue(isValidHcl("#abc123")))
# println(assertFalse(isValidHcl("#ab")))

# println("\nPid")
# println(assertTrue(isValidPid("123456789")))
# println(assertFalse(isValidPid("1234567890")))
# println(assertFalse(isValidPid("1A3456789")))
# println(assertFalse(isValidPid("1A345")))

# println("\nEcl")
# println(assertTrue(isValidEcl("brn")))
# println(assertFalse(isValidEcl("bla")))
# println(assertFalse(isValidEcl("asdf")))'

@test isValidEcl("asdf") == false


text:
function test_structure(sxt::SXTree, pairs...)
    function count_children(sxt, seq)
        curr = sxt[seq[1]]
        for i in seq[2:end]
            curr = curr[i]
        end
        SimplexTrees.num_children(curr)
    end

    all(count_children(sxt, seq) == n for (seq, n) in pairs)
end

function iter_count(f, n)
    i = 0
    for n in f(n)
        i += 1
    end
    i
end

function iter_collect(f, n)
    res = Vector{}[]
    for n in f(n)
        push!(res, vec(n))
    end
    res
end

@testset "Simplex insertion" begin
    @testset "Single simplex" begin

        # Check the structure of a simplex.
        # Build a 4-simplex and test the structure.
        structure4(sxt) = test_structure(sxt,
                                         [1] => 3,
                                         [1, 2] => 2,
                                         [1, 2, 3] => 1,
                                         [1, 2, 3, 4] => 0,
                                         [1, 2, 4] => 0,
                                         [1, 3] => 1,
                                         [1, 3, 4] => 0,

                                         [2] => 2,
                                         [2, 3] => 1,
                                         [2, 3, 4] => 0,
                                         [2, 4] => 0,

                                         [3] => 1,
                                         [3, 4] => 0,

                                         [4] => 0)
        sxt = SXTree(4)
        insert_simplex!(sxt, collect(1:4))
        @test structure4(sxt)

        # Build a full 6-simplex and roughly test the structure.
        structure6(sxt) = test_structure(sxt,
                                         Pair.([1:i for i in 1:6], 5:-1:0)...,
                                         Pair.([2:i for i in 2:6], 4:-1:0)...,
                                         Pair.([3:i for i in 3:6], 3:-1:0)...,
                                         Pair.([4:i for i in 4:6], 2:-1:0)...,
                                         Pair.([4:i for i in 5:6], 1:-1:0)...,
                                         [6] => 0)

        sxt = SXTree(6)
        insert_simplex!(sxt, collect(1:6))
        @test structure6(sxt)

    end

    @testset "Case from article" begin

        # Check the structure of a complex.
        # See: https://hal.archives-ouvertes.fr/file/index/docid/707901/filename/RR-7993.pdf
        # page 5, first connected component for the image.
        article_structure(sxt) = test_structure(sxt,
                                                [1] => 2,
                                                [1, 2] => 1,
                                                [1, 2, 3] => 0,
                                                [1, 3] => 0,

                                                [2] => 3,
                                                [2, 3] => 2,
                                                [2, 3, 4] => 1,
                                                [2, 3, 4, 5] => 0,
                                                [2, 4] => 1,
                                                [2, 4, 5] => 0,
                                                [2, 5] => 0,

                                                [3] => 2,
                                                [3, 4] => 1,
                                                [3, 4, 5] => 0,
                                                [3, 5] => 0,

                                                [4] => 1,
                                                [4, 5] => 0,

                                                [5] => 0)

        # Insert some simplices and make sure all of them were inserted properly.
        sxt = SXTree(5)
        # Sanity check.
        @test n_vertices(sxt) == 5

        insert_simplex!(sxt, [1, 3, 2])
        insert_simplex!(sxt, [2, 3, 4, 5])

        @test article_structure(sxt)
        @test dim(sxt) == 3

        # Insert in a different order
        sxt = SXTree(5)
        insert_simplex!(sxt, [1, 2])
        insert_simplex!(sxt, [5, 4])
        insert_simplex!(sxt, [3, 2])
        insert_simplex!(sxt, [3, 5])
        insert_simplex!(sxt, [1, 3])
        insert_simplex!(sxt, [2, 1, 3])
        insert_simplex!(sxt, [4, 5, 3, 2])

        @test article_structure(sxt)
        @test dim(sxt) == 3

        # Add some simplices that were already added and check again.
        insert_simplex!(sxt, [3, 4])
        insert_simplex!(sxt, [3, 4])
        insert_simplex!(sxt, [3, 4])
        insert_simplex!(sxt, [3, 4])

        @test article_structure(sxt)

        # Invalid arguments.
        @test_throws ArgumentError insert_simplex!(sxt, [1, 5, 6])
        @test_throws ArgumentError insert_simplex!(sxt, [1, 5, 2, 3, 1])

        # Nothing broke when attempting to call with invalid arguments.
        @test article_structure(sxt)
        @test dim(sxt) == 3
    end
end

@testset "Iterators" begin

    @testset "Cousins" begin
        cousin_count(sxt) = iter_count(cousins, sxt)
        # One element list.
        sxt = SXTree(2)
        insert_simplex!(sxt, [1, 2])

        @test first(cousins(sxt[1])) == sxt[1]
        @test cousin_count(sxt[1][2]) == 1

        # Build the test complex.
        sxt = SXTree(5)
        insert_simplex!(sxt, [1, 3, 2])
        insert_simplex!(sxt, [2, 3, 4, 5])

        @test first(cousins(sxt[1])) == sxt[1]

        @test cousin_count(sxt[2, 5]) == cousin_count(sxt[4][5]) == 3
        @test cousin_count(sxt[2, 3, 5]) == cousin_count(sxt[3][4][5]) == 3
    end

    @testset "Upwards" begin
        up_count(sxt) = iter_count(ancestors, sxt)

        sxt = SXTree(6)
        insert_simplex!(sxt, collect(1:6))

        @test up_count(sxt[6]) == 1
        @test up_count(sxt[1, 2, 3, 4, 5, 6]) == 6

        @test label.(collect(ancestors(sxt[4, 5, 6]))) == [6, 5, 4]

        @test length(ancestors(sxt[1, 2, 3, 4, 5, 6])) == 6
    end

    @testset "Vector representation" begin
        sxt = SXTree(6)
        insert_simplex!(sxt, collect(1:6))

        @test vec(sxt[1, 2]) == [1, 2]
        @test vec(sxt[2, 5, 6]) == [2, 5, 6]
        @test vec(sxt[1, 2, 3, 4, 5, 6]) == [1, 2, 3, 4, 5, 6]
    end

    @testset "BF iteration" begin
        sxt = SXTree(3)

        succ_bf_count(sxt) = iter_count(successors, sxt)
        succ_bf_collect(sxt) = iter_collect(successors, sxt)

        @test succ_bf_count(sxt) == 3
        @test succ_bf_count(sxt[3]) == 0
        insert_simplex!(sxt, collect(1:3))
        @test succ_bf_count(sxt) == 7
        @test succ_bf_collect(sxt) == [[1], [2], [3], [1,2], [1,3], [2,3], [1,2,3]]
        @test succ_bf_collect(sxt[1]) == [[1,2], [1,3], [1,2,3]]
    end

    @testset "DF iteration" begin
        sxt = SXTree(3)

        succ_df_count(sxt) = iter_count(x -> successors(x, DFIter), sxt)
        succ_df_collect(sxt) = iter_collect(x -> successors(x, DFIter), sxt)

        @test succ_df_count(sxt) == 3
        @test succ_df_count(sxt[3]) == 0
        insert_simplex!(sxt, collect(1:3))
        @test succ_df_count(sxt) == 7
        @test succ_df_collect(sxt) == [[1], [1,2], [1,2,3], [1,3], [2], [2,3], [3]]
        @test succ_df_collect(sxt[1]) == [[1,2], [1,2,3], [1,3]]
    end
end

@testset "Cofaces" begin
    sxt6 = SXTree(6)
    insert_simplex!(sxt6, collect(1:6))

    sxt4 = SXTree(4)
    insert_simplex!(sxt4, collect(1:4))

    @testset "is_face" begin
        @test is_face([1, 2, 3], sxt6[1, 2, 3, 4, 5, 6])
        @test is_face([4, 2, 6], sxt6[1, 2, 3, 4, 5, 6])
        @test !is_face([4, 2, 6], sxt6[1, 2, 3, 4, 5])

        @test_throws ArgumentError !is_face([1, 1], sxt6[1, 2, 3, 4, 5, 6])
    end

    @testset "cofaces" begin
        @test length(cofaces(sxt6, [1, 2, 3, 4, 5])) == 2
        @test length(cofaces(sxt6, [2, 3, 4, 5])) == 4
        @test length(cofaces(sxt6, [2, 4, 5])) == 8

        cfcs = vec.(cofaces(sxt4, [1, 2]))
        @test length(cfcs) == 4
        @test [1, 2] in cfcs &&
            [1, 2, 4] in cfcs &&
            [1, 2, 3] in cfcs &&
            [1, 2, 3, 4] in cfcs
    end
end

@testset "Simplex removal" begin
    @testset "rem_node!" begin
        sxt6 = SXTree(6)
        insert_simplex!(sxt6, collect(1:6))

        SimplexTrees.rem_node!(sxt6, sxt6[1,2,3,4,5,6])
        @test iter_count(successors, sxt6[1,2,3,4,5]) == 0

        sxt4 = SXTree(4)
        insert_simplex!(sxt4, collect(1:4))
        SimplexTrees.rem_node!(sxt4, sxt4[1, 2])
        @test iter_count(successors, sxt4[1]) == 3

        @test_throws ArgumentError SimplexTrees.rem_node!(sxt4, sxt4[1])
    end

    @testset "rem_simplex!" begin
        sxt6 = SXTree(6)
        insert_simplex!(sxt6, collect(1:6))

        rem_simplex!(sxt6, [1,2,3,4,5,6])
        @test dim(sxt6) == 4

        insert_simplex!(sxt6, collect(1:6))
        rem_simplex!(sxt6, sxt6[1,2,3,4,5,6])
        @test dim(sxt6) == 4

        rem_simplex!(sxt6, [1,2])
        @test dim(sxt6) == 4

        rem_simplex!(sxt6, [2,3])
        rem_simplex!(sxt6, [3,4])
        rem_simplex!(sxt6, [4,5])
        rem_simplex!(sxt6, [5,6])
        @test dim(sxt6) == 2
        @test n_simplices(sxt6) == 20
    end
end


text:
using Test, Dates, Julianda

foo = Julianda.Config.loadConfig("../config")

# Checks for the right response depending on the time
#TODO: redo pricing tests

max_stars_repo_name: CalebDepatie/Julianda

max_forks_repo_name: CalebDepatie/Julianda


text:
export edge_connectivity, min_edge_cut
export connectivity, min_cut

# min_cut(G,s,t) -- min cut separating s and t
# min_cut(G) -- min cut separating the graph
# flag == true means we want s/t cut

"""
`min_cut(G)` returns a minimum size set of vertices that disconnects `G`.

`min_cut(G,s,t)` returns a minimum size cut that separates `s` and `t`.
"""
function min_cut(G::SimpleGraph{T},s::T, t::T,flag::Bool=true)::Set{T} where T
    n = NV(G)
    m = NE(G)
    n*(n-1)!=2m || error("Graph must not be complete")

    if flag
        has(G,s)    || error("$s is not a vertex of this graph")
        has(G,t)    || error("$t is not a vertex of this graph")
        s!=t        || error("two vertices must be different")
        !has(G,s,t) || error("vertices $s and $t cannot be adjacent")
    else
        if cache_check(G,:min_cut)
            return cache_recall(G,:min_cut)
        end

        if !is_connected(G)
            X = Set{T}()
            cache_save(G,:min_cut,X)
            return X
        end
    end

    VV = vlist(G)
    EE = elist(G)

    MOD = Model(get_solver())

    @variable(MOD,a[VV],Bin)  # in part 1
    @variable(MOD,b[VV],Bin)  # in part 2
    @variable(MOD,c[VV],Bin)  # in cut set

    for v in VV
        @constraint(MOD, a[v]+b[v]+c[v]==1)
    end

    if flag
        @constraint(MOD,a[s]==1)
        @constraint(MOD,b[t]==1)
    end

    @constraint(MOD, sum(a[v] for v in VV)>=1)
    @constraint(MOD, sum(b[v] for v in VV)>=1)

    for e in EE
        v,w = e
        @constraint(MOD,a[v]+b[w]<=1)
        @constraint(MOD,a[w]+b[v]<=1)
    end

    @objective(MOD, Min, sum(c[v] for v in VV))

    optimize!(MOD)
    status = Int(termination_status(MOD))

    C =  value.(c)
    X = Set(v for v in VV if C[v] > 0.5)

    if !flag
        cache_save(G,:min_cut,X)
    end
    return X
end

function min_cut(G::SimpleGraph)::Set
    n = NV(G)
    n>1 || error("Graph must have at least two vertices")
    s = first(G.V)
    return min_cut(G,s,s,false)
end

"""
`connectivity(G)` returns the (vertex) connectivity of `G`, i.e.,
the minimum size of a vertex cut set. If `G` is a complete graph with
`n` vertices, the connectivity is `n-1` (or `0` for the empty graph).
"""
function connectivity(G::SimpleGraph)::Int
    n = NV(G)
    m = NE(G)
    if n*(n-1) == 2m   # graph is complete
        return max(n-1,0)  # what is the connectivity of K_0?
    end

    return length(min_cut(G))
end

# min_edge_cut(G,s,t) -- min edge cut separating s and t
# min_edge_cut(G) -- min edge cut separating the graph
# flag == true means we want s/t cut

"""
`min_edge_cut(G)` returns a minimum size set of edges whose removal
disconnects `G`. The graph must have at least two vertices.

`min_edge_cut(G,s,t)` is a minimum size set of edges whose removal
separates vertices `s` and `t`.
"""
function min_edge_cut(G::SimpleGraph{T}, s::T, t::T, flag::Bool=true)::Set{Tuple{T,T}} where T
    n = NV(G)
    n > 1 || error("Graph must have at least two vertices")

    if flag
        has(G,s) || error("$s is not a vertex of this graph")
        has(G,t) || error("$t is not a vertex of this graph")
        s!=t     || error("vertices must not be the same")
    else
        if cache_check(G,:min_edge_cut)
            return cache_recall(G,:min_edge_cut)
        end

        if !is_connected(G)
            X = Set{Tuple{T,T}}()
            cache_save(G,:min_edge_cut,X)
            return X
        end
    end

    VV = vlist(G)
    EE = elist(G)

    MOD = Model(get_solver())

    @variable(MOD,a[VV],Bin)  # in part 1
    @variable(MOD,b[VV],Bin)  # in part 2
    @variable(MOD,c[EE],Bin)  # span between the parts

    for v in VV
        @constraint(MOD,a[v]+b[v]==1)
    end

    @constraint(MOD,sum(a[v] for v in VV) >= 1)
    @constraint(MOD,sum(b[v] for v in VV) >= 1)

    if flag
        @constraint(MOD,a[s]==1)
        @constraint(MOD,b[t]==1)
    end

    for e in EE
        u,v = e
        @constraint(MOD,a[u]+b[v]-1 <= c[e])
        @constraint(MOD,a[v]+b[u]-1 <= c[e])
    end

    @objective(MOD, Min, sum(c[e] for e in EE))

    optimize!(MOD)
    status = Int(termination_status(MOD))

    C =  value.(c)

    X = Set(e for e in EE if C[e] > 0.5)

    if !flag
        cache_save(G,:min_edge_cut,X)
    end
    return X
end

function min_edge_cut(G::SimpleGraph)::Set
    n = NV(G)
    n>1 || error("Graph must have at least two vertices")
    s = first(G.V)
    return min_edge_cut(G,s,s,false)
end

"""
`edge_connectivity(G)` returns the size of a minimum edge cut of `G`.

`edge_connectivity(G,s,t)` determines the minimum size of an edge cut
separating `s` and `t`.
"""
edge_connectivity(G::SimpleGraph)::Int = length(min_edge_cut(G))
edge_connectivity(G::SimpleGraph,s,v) = length(min_edge_cut(G,s,t))


text:
length_local(x::Union{<:PencilArray, <:GlobalPencilArray}) = length(parent(x))

"""
    size(x::PencilArray)

Return global dimensions of a `PencilArray` in logical order.

Defined as `size_global(x, LogicalOrder())`.
"""
Base.size(x::Union{<:PencilArray, <:GlobalPencilArray}) =
    size_global(x, LogicalOrder())

"""
    length(x::PencilArray)

Get the *global* number of elements stored in the `PencilArray`.
"""
Base.length(x::Union{<:PencilArray, <:GlobalPencilArray}) = prod(size(x))

# By default, eachindex calls length(x) in this case.
# Since length returns the *global* length, we override it to return an iterator
# over the *local* length.
Base.eachindex(::IndexLinear, x::Union{<:PencilArray, <:GlobalPencilArray}) =
    Base.OneTo(length_local(x))

"""
    size_local(x::PencilArray, [order = LogicalOrder()])
    size_local(x::PencilArrayCollection, [order = LogicalOrder()])

Local dimensions of the data held by the `PencilArray`.

See also [`size_local(::Pencil)`](@ref).
"""
size_local(x::MaybePencilArrayCollection, args...; kwargs...) =
    (size_local(pencil(x), args...; kwargs...)..., extra_dims(x)...)

size_local(x::GlobalPencilArray, args...; kwargs...) =
    size_local(parent(x), args...; kwargs...)

"""
    size_global(x::PencilArray, [order = LogicalOrder()])
    size_global(x::PencilArrayCollection, [order = LogicalOrder()])

Global dimensions associated to the given array.

By default, the logical dimensions of the dataset are returned.

If `order = LogicalOrder()`, this is the same as `size(x)`.

See also [`size_global(::Pencil)`](@ref).
"""
size_global(x::MaybePencilArrayCollection, args...; kw...) =
    (size_global(pencil(x), args...; kw...)..., extra_dims(x)...)

size_global(x::GlobalPencilArray, args...; kwargs...) =
    size_global(parent(x), args...; kwargs...)