
"""
Function wrappers to package a dataset for use.

Adding your own dataset:
A dataset should be wrapped as a `dataset_name()` function that returns a dictionary containing the key/value pairs:
	- dataset_name :: (string) identifier for dataset
	- X_train :: (N_tr, N_x) tensor of type float32 where N_tr is the number of data points and N_x is the input dimensionality
	- Y_train :: for regression, a (N_tr, N_y) tensor of type float32 where N_y is the output dimensionality (typically 1)
			  :: for (binary) classification, a (N_tr, ) tensor of type int64 representing the output classes as integers (0-indexed)
	- X_train_min :: a list of length N_x, representing the lower bound of the training distribution for each input dimension
	- X_train_max :: a list of length N_x, representing the upper bound of the training distribution for each input dimension
	- (optional) X_test :: test set analogous to Y_train
	- (optional) Y_test :: test set analogous to Y_train
	- (optional) X_train_mean :: means of each feature of the training set
	- (optional) X_train_std :: standard deviations of each feature of the training set
	- ... others as necessary
"""

import numpy as np
import pandas as pd
import torch
from sklearn.model_selection import train_test_split
from sklearn.utils import resample


def toy1():
	""" 1D regression. """
	def f(x):
		return (-1 * (x ** 4)) + (3 * (x ** 2)) + 1

	X_train = torch.tensor([-2, -1.8, -1, 1, 1.8, 2]).unsqueeze(dim=1)
	Y_train = f(X_train)

	X_train_min = [-3.0]
	X_train_max = [3.0]
	return {"dataset_name": "toy1", "X_train": X_train, "Y_train": Y_train, "X_train_min": X_train_min, "X_train_max": X_train_max}


def toy2():
	""" 2D classification. 3 classes, with 8 data points for each class. Data was generated by a mixture of Gaussians. """
	
	X_train = torch.tensor([[-0.5568, -2.8603],
							[-2.8323,  0.9413],
							[ 2.0402,  2.6044],
							[-0.3121, -3.2628],
							[-3.2933,  0.8421],
							[ 2.2045,  2.6692],
							[-0.0165, -3.1644],
							[-2.9325,  0.8448],
							[ 2.3423,  2.8354],
							[-0.2595, -2.9736],
							[-2.9650,  0.7760],
							[ 2.2982,  2.6264],
							[-0.0605, -2.8752],
							[-2.4977,  1.2420],
							[ 1.8175,  3.1518],
							[-0.2743, -3.0461],
							[-2.7137,  1.7121],
							[ 1.9846,  3.5413],
							[ 0.3204, -2.8159],
							[-2.6090,  1.1718],
							[ 2.0794,  3.1475],
							[ 0.1280, -2.9005],
							[-2.9434,  1.1424],
							[ 1.9929,  2.9455]])
	Y_train = torch.tensor(range(3)).repeat(8)

	X_train_min = [-3.0, -3.0]
	X_train_max = [3.0, 3.0]	
	return {"dataset_name": "toy2", "X_train": X_train, "Y_train": Y_train, "X_train_min": X_train_min, "X_train_max": X_train_max}


def toy3():
	""" 1D regression. """
	X_train = torch.tensor([-7, -5, 2, 2.5]).unsqueeze(dim=1)
	Y_train = torch.tensor([-0.5, -1.0, 1.0, 0.7]).unsqueeze(dim=1)
	
	X_train_min = [-7.5]
	X_train_max = [3.0]
	return {"dataset_name": "toy3", "X_train": X_train, "Y_train": Y_train, "X_train_min": X_train_min, "X_train_max": X_train_max}


def toy4():
	""" 1D regression. """
	X_train = torch.tensor([-3.5, -2, 2, 3.5]).unsqueeze(dim=1)
	Y_train = torch.tensor([-1, 0.25, 3.5, 4.5]).unsqueeze(dim=1)

	X_train_min = [-4.0]
	X_train_max = [4.0]
	return {"dataset_name": "toy4", "X_train": X_train, "Y_train": Y_train, "X_train_min": X_train_min, "X_train_max": X_train_max}


def toy5():
	""" 2D binary classification. X_1 is a binary feature, and X_2 is continuous in [0, 1]. """
	X_train = torch.stack((torch.arange(2.0).repeat_interleave(10), torch.arange(0.1, 1.1, 0.1).repeat(2))).t()
	Y_train = torch.tensor([0] * 2 + [1] * 8 + [0] * 8 + [1] * 2)

	X_train_min = [0.0, 0.0]
	X_train_max = [1.0, 1.0]
	return {"dataset_name": "toy5", "X_train": X_train, "Y_train": Y_train, "X_train_min": X_train_min, "X_train_max": X_train_max}


def toy6():
	""" 1D regression. Outputs are intentionally noisy. """
	
	def f(x):
		return 5 * np.cos(x / 1.7)

	X_train = torch.tensor([-2.7, -2.3, -2, 2, 2.3, 2.7]).unsqueeze(1)
	Y_train = f(X_train) + torch.tensor([0.4, 0.2, 0.7, -2.5, 0.6, 2.1]).unsqueeze(1)
	
	X_train_min = [-3.0]
	X_train_max = [3.0]
	return {"dataset_name": "toy6", "X_train": X_train, "Y_train": Y_train, "X_train_min": X_train_min, "X_train_max": X_train_max}


def compas_dataset(csv_xfilename, csv_yfilename, with_race=True):
	""" COMPAS dataset from 2016 ProPublica study:
			Jeff Larson, Surya Mattu, Lauren Kirchner, and Julia Angwin. How We Analyzed the COMPAS
			Recidivism Algorithm. ProPublica, 2016.
		9D binary classification (8D without race as an explicit feature).
	"""
	X = pd.read_csv(csv_xfilename)
	Y = np.loadtxt(csv_yfilename)
	X = X.values

	X_train = torch.tensor(X[:, 1:]).float()
	Y_train = torch.tensor(Y).long()
	X_train_race = X_train[:, -1]
	if not with_race:
		X_train = X_train[:, :-1]

	# Means and standard deviation for each feature.
	X_train_mean, X_train_std = [], []
	for i in range(X_train.shape[1]):
		X_train_mean.append(X_train[:,i].mean().item())
		X_train_std.append(X_train[:,i].std().item())

	# Standardization.
	for j in [0,2,3]:
		X_train[:,j] = (X_train[:,j] - X_train_mean[j]) / X_train_std[j]

	# Min/max.
	X_train_min, X_train_max = [], []
	for i in range(X_train.shape[1]):
		X_train_min.append(X_train[:,i].min().item())
		X_train_max.append(X_train[:,i].max().item())

	return {"dataset_name": "compas", "X_train": X_train, "Y_train": Y_train,
			"X_train_min": X_train_min, "X_train_max": X_train_max, "X_train_mean": X_train_mean, 
			"X_train_std": X_train_std, "X_train_race": X_train_race}


def credit_dataset(csv_filename):
	""" Give Me Some Credit dataset, taken from: http://www.kaggle.com/c/GiveMeSomeCredit/.
		10D binary classification.
	"""
	data = pd.read_csv(csv_filename)
	data.dropna(inplace=True)
	data = data.values[:,1:]

	# Filterig outliers.
	for i in [3,4,5]:
		data = data[(data[:,i] <= np.quantile(data[:,i], 0.95))]
	for i in [1,7,8,9,10]:
		data = data[(data[:,i] <= np.quantile(data[:,i], 0.998))]
	
	data_train, data_test = train_test_split(data, test_size=0.1, random_state=13)
	X_test = data_test[:,1:]
	Y_test = data_test[:,0]

	# Means and standard deviation for each feature.
	X_train_mean, X_train_std = [], []
	for i in range(1, 11):
		X_train_mean.append(data_train[:,i].mean())
		X_train_std.append(data_train[:,i].std())

	# Standardization.
	for j in [1,4,5]:
		X_test[:,j] = (X_test[:,j] - X_train_mean[j]) / X_train_std[j]
		data_train[:,j+1] = (data_train[:,j+1] - X_train_mean[j]) / X_train_std[j]
	
	# Upsampling minority class (Y = 1).
	majority = data_train[data_train[:,0] == 0]
	minority = data_train[data_train[:,0] == 1]
	minority_upsampled = resample(minority, replace=True, n_samples=round(len(majority)*0.5), random_state=13)
	data_train = np.concatenate((majority, minority_upsampled))
	np.random.shuffle(data_train)
	X_train = data_train[:,1:]
	Y_train = data_train[:,0]

	# Conversion to tensors.
	X_train = torch.tensor(X_train).float()
	X_test = torch.tensor(X_test).float()
	Y_train = torch.tensor(Y_train).long()
	Y_test = torch.tensor(Y_test).long()

	# "Blind dataset" lacking individuals with age < 35
	X_trainbl = X_train.clone()
	Y_trainbl = Y_train.clone()
	Y_trainbl = Y_trainbl[(X_trainbl[:,1] * X_train_std[1] + X_train_mean[1] >= 35).nonzero().squeeze(dim=1)]
	X_trainbl = X_trainbl[(X_trainbl[:,1] * X_train_std[1] + X_train_mean[1] >= 35).nonzero().squeeze(dim=1)]

	# Min/max.
	X_train_min, X_train_max = [], []
	for i in range(10):
		X_train_min.append(X_trainbl[:,i].min().item())
		X_train_max.append(X_trainbl[:,i].max().item())

	return {"dataset_name": "give_me_some_credit", "X_train": X_train, "Y_train": Y_train,
			"X_train_min": X_train_min, "X_train_max": X_train_max, "X_test": X_test, "Y_test": Y_test,
			"X_train_mean": X_train_mean, "X_train_std": X_train_std, "X_train_blind": X_trainbl, "Y_train_blind": Y_trainbl}
