import os
import cv2
import random
import numpy as np
from tensorflow.keras.utils import to_categorical
from scripts.consts import class_dict


def get_data(path, split=0.2):

    X, y = [], []

    for directory in os.listdir(path):

        dirpath = os.path.join(path, directory)
        print(directory, len(os.listdir(dirpath)))

        for file in os.listdir(dirpath):

            filepath = os.path.join(dirpath, file)
            img = cv2.imread(filepath, cv2.IMREAD_UNCHANGED)
            
            if img.shape != (360, 363, 3):
                img = cv2.resize(img, (360, 363), cv2.INTER_CUBIC)

            X.append(img)
            y.append(class_dict[directory])

    data = list(zip(X, y))

    random.shuffle(data)

    X, y = zip(*data)

    num_train = int((1.0 - split) * len(y))

    X_train, X_valid = np.array(X[:num_train]).astype(
        'float32'), np.array(X[num_train:]).astype('float32')
    y_train, y_valid = np.array(
        y[:num_train]).reshape(-1, 1), np.array(y[num_train:]).reshape((-1, 1))

    X_train = X_train / 255.0
    X_valid = X_valid / 255.0

    y_train, y_valid = to_categorical(y_train), to_categorical(y_valid)

    print(X_train.shape, y_train.shape)
    print(X_valid.shape, y_valid.shape)

    return X_train, y_train, X_valid, y_valid

[(13, 'os.listdir', 'os.listdir', False, 'import os\n'), (31, 'random.shuffle', 'random.shuffle', False, 'import random\n'), (15, 'os.path.join', 'os.path.join', False, 'import os\n'), (18, 'os.listdir', 'os.listdir', False, 'import os\n'), (45, 'tensorflow.keras.utils.to_categorical', 'to_categorical', False, 'from tensorflow.keras.utils import to_categorical\n'), (45, 'tensorflow.keras.utils.to_categorical', 'to_categorical', False, 'from tensorflow.keras.utils import to_categorical\n'), (20, 'os.path.join', 'os.path.join', False, 'import os\n'), (21, 'cv2.imread', 'cv2.imread', False, 'import cv2\n'), (16, 'os.listdir', 'os.listdir', False, 'import os\n'), (24, 'cv2.resize', 'cv2.resize', False, 'import cv2\n'), (37, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (38, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (39, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (40, 'numpy.array', 'np.array', True, 'import numpy as np\n')]