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, 28, 26, 67), 'os.path.join', 'os.path.join', ({(26, 41, 26, 45): 'path', (26, 47, 26, 66): 'const.FILE_METADATA'}, {}), '(path, const.FILE_METADATA)', False, 'import os\n'), ((27, 28, 27, 62), 'os.path.join', 'os.path.join', ({(27, 41, 27, 45): 'path', (27, 47, 27, 61): 'const.DIR_TEST'}, {}), '(path, const.DIR_TEST)', False, 'import os\n'), ((28, 29, 28, 64), 'os.path.join', 'os.path.join', ({(28, 42, 28, 46): 'path', (28, 48, 28, 63): 'const.DIR_TRAIN'}, {}), '(path, const.DIR_TRAIN)', False, 'import os\n'), ((32, 23, 32, 58), 'os.path.join', 'os.path.join', ({(32, 36, 32, 40): 'path', (32, 42, 32, 57): 'const.DIR_TRAIN'}, {}), '(path, const.DIR_TRAIN)', False, 'import os\n'), ((36, 22, 36, 56), 'os.path.join', 'os.path.join', ({(36, 35, 36, 39): 'path', (36, 41, 36, 55): 'const.DIR_TEST'}, {}), '(path, const.DIR_TEST)', False, 'import os\n'), ((39, 13, 39, 58), 'pandas.read_csv', 'pd.read_csv', (), '', True, 'import pandas as pd\n'), ((77, 17, 77, 36), 'numpy.asarray', 'np.asarray', ({(77, 28, 77, 35): 'x_train'}, {}), '(x_train)', True, 'import numpy as np\n'), ((78, 17, 78, 36), 'numpy.asarray', 'np.asarray', ({(78, 28, 78, 35): 'y_train'}, {}), '(y_train)', True, 'import numpy as np\n'), ((79, 16, 79, 34), 'numpy.asarray', 'np.asarray', ({(79, 27, 79, 33): 'x_test'}, {}), '(x_test)', True, 'import numpy as np\n'), ((80, 16, 80, 34), 'numpy.asarray', 'np.asarray', ({(80, 27, 80, 33): 'y_test'}, {}), '(y_test)', True, 'import numpy as np\n'), ((55, 27, 55, 75), 'os.path.join', 'os.path.join', ({(55, 40, 55, 58): 'train_img_dir_path', (55, 60, 55, 74): 'image_row[row]'}, {}), '(train_img_dir_path, image_row[row])', False, 'import os\n'), ((57, 24, 57, 66), 'cv2.imread', 'cv2.imread', ({(57, 35, 57, 43): 'img_path', (57, 45, 57, 65): 'cv2.IMREAD_GRAYSCALE'}, {}), '(img_path, cv2.IMREAD_GRAYSCALE)', False, 'import cv2\n'), ((59, 24, 59, 53), 'cv2.resize', 'cv2.resize', ({(59, 35, 59, 40): 'image', (59, 42, 59, 52): '(128, 128)'}, {}), '(image, (128, 128))', False, 'import cv2\n'), ((67, 27, 67, 74), 'os.path.join', 'os.path.join', ({(67, 40, 67, 57): 'test_img_dir_path', (67, 59, 67, 73): 'image_row[row]'}, {}), '(test_img_dir_path, image_row[row])', False, 'import os\n'), ((69, 24, 69, 66), 'cv2.imread', 'cv2.imread', ({(69, 35, 69, 43): 'img_path', (69, 45, 69, 65): 'cv2.IMREAD_GRAYSCALE'}, {}), '(img_path, cv2.IMREAD_GRAYSCALE)', False, 'import cv2\n'), ((71, 24, 71, 53), 'cv2.resize', 'cv2.resize', ({(71, 35, 71, 40): 'image', (71, 42, 71, 52): '(128, 128)'}, {}), '(image, (128, 128))', False, 'import cv2\n')]