
import numpy as np
from src.plotting import plot_stuff

# Custom's selection
exp = 'exp_synthetic_2_clusters'  # Figure 1 left
# exp = 'exp_synthetic_6_clusters'  # Figure 1 right
# exp = 'exp_real_lenk'  # Figure 2 top-left
# exp = 'exp_real_movies'  # Figure 2 top-right
# exp = 'exp_real_jester'  # Figure 2 bottom

if exp == 'exp_synthetic_2_clusters':
    methods = ['ITL', 'unconditional', 'conditional']
    saved_file_name = 'saved_results/synth_2_clusters_results.npy'
    dataset = 'synthetic-regression-feature-2-CLUSTERS'
elif exp == 'exp_synthetic_6_clusters':
    methods = ['ITL', 'unconditional', 'conditional']
    saved_file_name = 'saved_results/synth_6_clusters_results.npy'
    dataset = 'synthetic-regression-feature-6-CLUSTERS'
elif exp == 'exp_real_lenk':
    methods = ['ITL', 'unconditional', 'conditional']
    saved_file_name = 'saved_results/lenk_results.npy'
    dataset = 'lenk'
elif exp == 'exp_real_movies':
    methods = ['ITL', 'unconditional', 'conditional']
    saved_file_name = 'saved_results/movies_results_2.npy'
    dataset = 'movies'
elif exp == 'exp_real_jester':
    methods = ['ITL', 'unconditional', 'conditional']
    saved_file_name = 'saved_results/jester_results_2.npy'
    dataset = 'jester'

results = np.load(saved_file_name, allow_pickle='TRUE').item()
plot_stuff(results, methods, dataset)