import numpy as np
import torch
from torch.autograd import Variable
from torch.utils.data import DataLoader, TensorDataset, random_split
# import random


#RandomSeed = 191
dim = 100
emb = 100
NumSample = 20
nz = 5
bigvariance = 1
smallvariance = 0.1
#noise = 0
noise = 0.02
TrainRatio = 0.5

def EmbDataSet(beta, RandomSeed, d=dim, emb = emb, N=NumSample, TrainRatio=TrainRatio):
    N_train = int(N * TrainRatio)
    N_test = N - N_train
    np.random.seed(RandomSeed)
    # Create x data for the train set
    #x_values = [np.concatenate((np.random.normal(0, bigvariance, size=(1, emb))[0], np.random.normal(0, smallvariance, size=(1,d-emb))[0])) for i in range(N_train)]
    x_values = [np.concatenate(
        (np.random.uniform(-bigvariance, bigvariance, size=(1, emb))[0], np.random.uniform(-smallvariance, smallvariance, size=(1, d - emb))[0]))
                for i in range(N_train)]
    #print(x_values)
    x_train = np.array(x_values, dtype=np.float32)
    # x_train = x_train.reshape(-1, 1)
    # print(x_train)
    # print(np.shape(x_train))

    # Create the d-dim model parameter beta
    #beta = np.zeros((d, 1))
    #beta[0:nz] = np.random.normal(0, 2, size=(nz, 1))
    #beta[0:nz] = np.ones(shape=(nz,1))
    beta = beta
    print(beta[0:nz])
    #beta[nz:dim] = -0
    #beta = np.random.normal(0, 1, size=(d, 1))
    #beta[emb:d] = 0

    # Create y data for the train set
    y_values = np.dot(x_train, beta)
    # Add noise to y
    y_values += np.random.normal(0, noise, size=(N_train, 1))
    y_train = np.array(y_values, dtype=np.float32)
    # print(type(y_train[0,0,0]))
    # print(y_train)

    # dataset = np.concatenate((x_train, y_train), axis=1)
    # print(dataset)

    # Create x data for the test set
    #x_values = [np.concatenate(
    #    (np.random.normal(0, bigvariance, size=(1, emb))[0], np.random.normal(0, smallvariance, size=(1, d - emb))[0]))
    #            for i in range(N_test)]
    x_values = [np.concatenate(
        (np.random.uniform(-bigvariance, bigvariance, size=(1, emb))[0], np.random.uniform(-smallvariance, smallvariance, size=(1, d - emb))[0]))
        for i in range(N_test)]
    x_test = np.array(x_values, dtype=np.float32)

    # Create y data for the test set
    y_values = np.dot(x_test, beta)
    # Add noise to y
    y_values += np.random.normal(0, noise, size=(N_test, 1))
    y_test = np.array(y_values, dtype=np.float32)

    return x_train, y_train, x_test, y_test, d, N_train, N_test
