import pandas as pd
import os
import numpy as np
import cv2
from utils import constants as const
import matplotlib.pyplot as plt

class DataLoader:

    def load_data():
        '''
        This function is handling the data loading and pre-processing

        :return: (xtrain, ytrain), (xtest, ytest)
        '''
        print('**** Read data into DAugmentor ****')

        x_train = []
        y_train = []
        x_test = []
        y_test = []

        # setting the path to metadata
        path = const.PATH

        metadata_csv_path = os.path.join(path, const.FILE_METADATA)
        test_img_dir_path = os.path.join(path, const.DIR_TEST)
        train_img_dir_path = os.path.join(path, const.DIR_TRAIN)
        print(metadata_csv_path)

        # setting the path to train data
        x_train_path = os.path.join(path, const.DIR_TRAIN)
        print(x_train_path)

        # setting the path to train data
        x_test_path = os.path.join(path, const.DIR_TEST)

        # reading meta data file as dataframe
        df = pd.read_csv(metadata_csv_path, delimiter=',')

        # dataset format:
        #   image_name
        #   label
        #   data_type

        data_type_row = df["data_type"].tolist()
        image_row = df["image_name"].tolist()
        label_row = df["label"].tolist()

        data_rows = len(data_type_row)

        for row in range(data_rows):
            if (data_type_row[row] == "TRAIN"):
                # setting the path of the current image
                img_path = os.path.join(train_img_dir_path, image_row[row])
                # reading image
                image = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
                # downscaling image to 28x28
                image = cv2.resize(image, (128, 128))
                x_train.append(image)
                print("Loaded: " + img_path)
                # extracting labels
                y_train.append(label_row[row])

            if (data_type_row[row] == "TEST"):
                # setting the path of the current image
                img_path = os.path.join(test_img_dir_path, image_row[row])
                # reading image
                image = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
                # downscaling image to 28x28
                image = cv2.resize(image, (128, 128))
                x_test.append(image)
                print("Loaded: " + img_path)
                # extracting labels
                y_test.append(label_row[row])

        xtrain = np.asarray(x_train)
        ytrain = np.asarray(y_train)
        xtest = np.asarray(x_test)
        ytest = np.asarray(y_test)

        print(x_train[0].shape)
        print(x_train[0].shape)

        print(xtrain[0].shape)
        print(x_test[0].shape)

        #(X_train, y_train), (X_test, y_test)
        return (xtrain, ytrain), (xtest, ytest)
[(26, 'os.path.join', 'os.path.join', (['path', 'const.FILE_METADATA'], {}), False, 'import os\n'), (27, 'os.path.join', 'os.path.join', (['path', 'const.DIR_TEST'], {}), False, 'import os\n'), (28, 'os.path.join', 'os.path.join', (['path', 'const.DIR_TRAIN'], {}), False, 'import os\n'), (32, 'os.path.join', 'os.path.join', (['path', 'const.DIR_TRAIN'], {}), False, 'import os\n'), (36, 'os.path.join', 'os.path.join', (['path', 'const.DIR_TEST'], {}), False, 'import os\n'), (39, 'pandas.read_csv', 'pd.read_csv', (['metadata_csv_path'], {'delimiter': '""","""'}), True, 'import pandas as pd\n'), (77, 'numpy.asarray', 'np.asarray', (['x_train'], {}), True, 'import numpy as np\n'), (78, 'numpy.asarray', 'np.asarray', (['y_train'], {}), True, 'import numpy as np\n'), (79, 'numpy.asarray', 'np.asarray', (['x_test'], {}), True, 'import numpy as np\n'), (80, 'numpy.asarray', 'np.asarray', (['y_test'], {}), True, 'import numpy as np\n'), (55, 'os.path.join', 'os.path.join', (['train_img_dir_path', 'image_row[row]'], {}), False, 'import os\n'), (57, 'cv2.imread', 'cv2.imread', (['img_path', 'cv2.IMREAD_GRAYSCALE'], {}), False, 'import cv2\n'), (59, 'cv2.resize', 'cv2.resize', (['image', '(128, 128)'], {}), False, 'import cv2\n'), (67, 'os.path.join', 'os.path.join', (['test_img_dir_path', 'image_row[row]'], {}), False, 'import os\n'), (69, 'cv2.imread', 'cv2.imread', (['img_path', 'cv2.IMREAD_GRAYSCALE'], {}), False, 'import cv2\n'), (71, 'cv2.resize', 'cv2.resize', (['image', '(128, 128)'], {}), False, 'import cv2\n')]