# this file generates the data for the real linear plots
import numpy as np
import cascading_alg as c_alg
import other_alg as o_alg
import time
from os.path import exists

start = time.time()  # to see how long the entire experiment takes
n, d, K, trials = 1000000, 20, 4, 5  # horizon, feature dimension, number of recommended items, and number of trials
# process movielens data if needed, printing error and exiting if the data isn't present
if not exists('data/movies.npz'):
    data_processed = o_alg.process_movielens()
    if not data_processed:
        print('error -- file data/movies.npz does not exist -- see README')
        exit()
data = np.load('data/movies.npz')  # load data, which exists if we've gotten here
W, genre_list, genre_mat = data['arr_0'], data['arr_1'], data['arr_2']  # extract the three arrays
m, L = W.shape  # number of users and number of items

# all genre experiment
regret, stats = np.zeros((2, n, trials)), np.zeros((4, trials))
print('starting all genres')
for trial_idx in range(trials):
    # split into train and test, train features and compute optimal action as described in appendix C, run two linear
    #   algorithms, save some stats, and print progress update
    start_trial = time.time()
    train_rows = np.random.choice(m, int(m / 2), False)
    test_rows = np.setdiff1d(range(m), train_rows)
    W_train, W_test = W[train_rows, :], W[test_rows, :]
    phi, theta = o_alg.real_features(W_train, d)
    opt = o_alg.real_optimal(W_test, K)
    regret[0, :, trial_idx] = c_alg.cascade_linucb(n, K, phi, [W_test, opt])
    regret[1, :, trial_idx] = c_alg.cascade_woful(n, K, phi, [W_test, opt])
    stats[0, trial_idx] = W.shape[1]  # number movies for this experiment
    stats[1, trial_idx] = np.mean(W_test)  # average click probability
    stats[2, trial_idx] = np.mean(np.max(W_train[:, opt], 1))  # click probability of opt
    stats[3, trial_idx] = np.linalg.norm(theta)  # norm of theta (see appendix C)
    print('trial ' + str(trial_idx + 1) + ' of ' + str(trials) + ' done, time = ' + str(time.time() - start_trial))
np.savez('results/real_lin_all', regret, stats)  # save results for later plotting

# per genre experiment
for genre_idx in range(len(genre_list)):
    regret, stats = np.zeros((2, n, trials)), np.zeros((4, trials))
    W_genre = W[:, np.where(genre_mat[:, genre_idx + 1])[0]]  # just the movies with this genre tag
    print('starting genre \'' + genre_list[genre_idx] + '\'')
    for trial_idx in range(trials):
        # split into train and test, train features and compute optimal action as described in appendix C, run two
        #   linear algorithms, save some stats, and print progress update
        start_trial = time.time()
        train_rows = np.random.choice(m, int(m / 2), False)
        test_rows = np.setdiff1d(range(m), train_rows)
        W_train, W_test = W_genre[train_rows, :], W_genre[test_rows, :]
        phi, theta = o_alg.real_features(W_train, d)
        opt = o_alg.real_optimal(W_test, K)
        regret[0, :, trial_idx] = c_alg.cascade_linucb(n, K, phi, [W_test, opt])
        regret[1, :, trial_idx] = c_alg.cascade_woful(n, K, phi, [W_test, opt])
        stats[0, trial_idx] = W_genre.shape[1]  # number movies for this experiment
        stats[1, trial_idx] = np.mean(W_test)  # average click probability
        stats[2, trial_idx] = np.mean(np.max(W_train[:, opt], 1))  # click probability of opt
        stats[3, trial_idx] = np.linalg.norm(theta)  # norm of theta (see appendix C)
        print('trial ' + str(trial_idx + 1) + ' of ' + str(trials) + ' done, time = ' + str(time.time() - start_trial))
    np.savez('results/real_lin_' + genre_list[genre_idx], regret, stats)  # save results for later plotting

print('total time = ' + str(time.time() - start))  # show how long the entire experiment took
