# -*- coding: utf-8 -*-
# @Author  : Yandai
# @Email   : <EMAIL>
# @File    : ctgan_base.py
# @Time    : 2021-7-20 13:50
# @Software: PyCharm

from ctgan import CTGANSynthesizer
from ctgan import load_demo,load_adult
import warnings
warnings.filterwarnings("ignore")
import matplotlib.pyplot as plt
import numpy as np
from sklearn.model_selection import train_test_split
import pandas as pd

data = load_adult()
data_columns = data.columns
X_ori=data.values[:,0:-1]
y_ori=data.values[:,-1]
X, X_test, y, y_test = train_test_split(X_ori, y_ori, test_size=0.2, random_state=0)
np.savetxt("Results/adult_balance/x.csv",X,delimiter=",")
np.savetxt("Results/adult_balance/y.csv",y,delimiter=",")
np.savetxt("Results/adult_balance/x_test.csv",X_test,delimiter=",")
np.savetxt("Results/adult_balance/y_test.csv",y_test,delimiter=",")

num_names = ["f1","f3","f5",'f11', 'f12', 'f13']
num_cols = [0,2,4,10,11,12]
X_maj = X[np.where(y==0)]
y_maj = y[np.where(y==0)]
X_min = X[np.where(y==1)]
y_min = y[np.where(y==1)]
for i in range(10):
    X = np.concatenate((X,X_min),axis=0)
    y = np.concatenate((y,y_min),axis=0)

data = pd.DataFrame(np.concatenate((X,y.reshape(-1,1)),axis=1))
data.columns = data_columns
for i in num_cols:
    plt.hist(data.values[:,i])
    plt.title(num_names[num_cols.index(i)])
    plt.show()

# Names of the columns that are discrete
discrete_columns = [
    'f2',
    'f4',
    'f6',
    'f7',
    'f8',
    'f9',
    'f10',
    'f14',
    'y'
]

def Sample(iter):
    ctgan = CTGANSynthesizer(epochs=10000)
    ctgan.fit(data, discrete_columns)
    samples = ctgan.sample(20000)
    np.savetxt("Results/adult_balance/samples_adult_{}.csv".format(iter),samples,delimiter=",")

if __name__ == '__main__':
    for i in range(20):
        Sample(i)


