import torch
import numpy as np
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
import pandas as pd
from sklearn.decomposition import PCA


class TfDataset(Dataset):
    def __init__(self, filepath, save_filepath=None):
        xy = np.loadtxt(filepath, delimiter=',', dtype=np.float32)
        self.len = xy.shape[0]  # shape(多少行，多少列)
        self.x_data = torch.from_numpy(xy[:, :-1])
        self.y_data = torch.from_numpy(xy[:, [-1]])

        # 使用 PCA 进行降维
        pca = PCA(n_components=12)
        reduced_data = pca.fit_transform(self.x_data.numpy())

        # 将降维后的数据设置为新的 x_data
        self.x_data = torch.from_numpy(reduced_data)

        # 将降维后的数据保存为 CSV 文件（如果提供了保存路径）
        if save_filepath:
            df_reduced = pd.DataFrame(reduced_data,
                                      columns=['PC1', 'PC2', 'PC1', 'PC2', 'PC1', 'PC2', 'PC1', 'PC2', 'PC1',
                                               'PC2', 'PC1', 'PC1'])
            df_reduced['label'] = self.y_data.numpy()
            df_reduced['dborgot'] = 0  # 添加一列值为0的列,代表为糖尿病数据集
            df_reduced.to_csv(save_filepath, index=False)

    def __getitem__(self, index):
        return self.x_data[index], self.y_data[index]

    def __len__(self):
        return self.len

    # 初始化数据集并保存降维后的数据为 CSV 文件


train_dataset = TfDataset('./diabetes_012_health_indicators_BRFSS2015.csv', save_filepath='./reduced_data.csv')

#
# train_dataset = tfDataset('../data/resample.csv')