import numpy as np
import os
import pickle
import matplotlib.pyplot as plt
import sys

np.set_printoptions(suppress=True)

datasets_uci = ['bostonHousing', 'concrete', 'energy', 'kin8nm', 'naval-propulsion-plant', 'power-plant', 'protein-tertiary-structure', 'wine-quality-red', 'yacht']
dataset_names = ['boston_housing', 'concrete', 'energy', 'kin8nm', 'naval', 'power', 'protein', 'wine', 'yacht']
input_dims = {'boston_housing': 13, 'concrete': 8, 'energy': 8, 'kin8nm': 8, 'power': 4, 'protein': 9, 'wine': 11, 'yacht': 6, 'naval': 16}

for i, dataset in enumerate(datasets_uci):
    data_directory = 'UCI_Datasets//' + dataset + '//data//'
    data = np.loadtxt(data_directory + 'data.txt')
    if dataset == 'naval-propulsion-plant':
        data = data[:,:-1] # regress the first of 2 targets
    split_directory = 'gap_datasets//' + dataset_names[i] + '_gap'
    os.mkdir(split_directory)

    for j in range(input_dims[dataset_names[i]]): # for each input dim, create a different split that has a gap in that dimension
        # get the split indices
        # look at the jth input
        jth_input = data[:,j]
        target = data[:,-1]
        sorting_indeces = np.argsort(jth_input)

        # remove 1/3 of data in the 'middle' as a test set
        num_data = target.size
        test_size = int(np.ceil(num_data/3))
        test_indices = sorting_indeces[int(num_data/2) - int(test_size/2):int(num_data/2) + int(test_size/2)]
        train_indices = np.concatenate((sorting_indeces[:int(num_data/2) - int(test_size/2)], sorting_indeces[int(num_data/2) + int(test_size/2):]))
        
        # extract train set
        train_set = data[train_indices,:]

        # extract test set
        test_set = data[test_indices,:]

        # normalise the train set
        train_mean = 1/(train_set.shape[0])*np.sum(train_set, axis=0)
        train_set_normalised = train_set - train_mean # broadcast
        
        train_sd = np.std(train_set_normalised, axis=0)
        for n in range(train_sd.size):
            if train_sd[n] < 1e-6:
                train_sd[n] = 1 # prevent division by zero if all the inputs are the same value, as in naval
        train_set_normalised = train_set_normalised/train_sd # broadcast

        # extract a validation set - 10% of the trainset
        val_size = int(train_indices.size/10)
        val_set_normalised = train_set_normalised[-val_size:,:] 
        train_set_normalised = train_set_normalised[0:-val_size,:]

        # pickle
        filename = split_directory + '//' + dataset_names[i] + str(j) + '.pkl'
        with open(filename, 'wb') as f:
            pickle.dump([train_set, train_set_normalised, val_set_normalised, test_set, train_mean, train_sd], f)

        # write details to a text file
        np.savetxt(split_directory + '//train_indices_' + str(j) + '.txt', train_indices.astype(int), fmt='%i')
        np.savetxt(split_directory + '//test_indices' + str(j) + '.txt', test_indices.astype(int), fmt='%i')

        file = open(split_directory + '//details.txt','w') 
        file.write('train set size : {} \n'.format(train_indices.shape[0]))
        file.write('test set size: {} \n'.format(test_set.shape[0]))
        file.close()
