"""
Handling dataset with PyTorch
"""
import torch
import numpy as np
import pandas as pd
from torch.utils.data import Dataset, DataLoader


class MyDataset(Dataset):
    """
    My dataset class
    Do tasks with small memory usage in __init__
    Define function __len__
    Read data from disk in __getitem__
    """
    def __init__(self, csv_file_path, data_path, data_row, data_column):
        """
        Minor processing of dataset related objects
        Read data info here
        """
        self.csv_file_path = csv_file_path
        self.data_path = data_path
        # rows and columns where data start, start at 0
        self.data_row = data_row
        self.data_column = data_column

        # normally, self.csv should be small in size, shouldn't contain data
        self.csv = pd.read_csv(csv_file_path, dtype={'name': str})
        data = pd.read_csv(self.data_path, nrows=5)

        self.N = self.csv.shape[0]
        self.d = data.shape[1] - data_column - 1  # -1 for y

        self.features_names = list(data.columns[data_column+1:])

        # Print a summary of read csv
        print(csv_file_path + ' is read with N: %i, d: %i.' % (self.N, self.d))
        # self.csv.info(memory_usage='deep')
        print('csv rows:')
        print(self.csv.iloc[:5, :])
        print('data rows:')
        print(data.iloc[:5, :])

    def __len__(self):
        """
        Return the number of samples in the dataset
        """
        return self.N

    def __getitem__(self, indice):
        """
        Read and return the data with given idx from the disk
        idx: in {0, ..., N-1}^?
        y = np.array(DataFrame.iloc[:, 2], np.float) to get columns
        """
        rows = pd.read_csv(self.data_path, header=None,
                           skiprows=lambda x: x != self.data_row + indice)
        x = torch.Tensor(np.array(rows.iloc[:, self.data_column+1:],
                                  np.float).squeeze())
        y = torch.Tensor(np.array(rows.iloc[:, self.data_column],
                                  np.float).squeeze())

        return x, y


if __name__ == '__main__':
    torch.manual_seed(1)
    dataset = MyDataset('dataset.csv', 'dataset.csv', 1, 1)
    dataloader = DataLoader(dataset, batch_size=6, shuffle=True, num_workers=4)

    for i, [x, y] in enumerate(dataloader):
        print(x)
        print(y)
