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', (['path'], {}), False, 'import os\n'), (31, 'random.shuffle', 'random.shuffle', (['data'], {}), False, 'import random\n'), (15, 'os.path.join', 'os.path.join', (['path', 'directory'], {}), False, 'import os\n'), (18, 'os.listdir', 'os.listdir', (['dirpath'], {}), False, 'import os\n'), (45, 'tensorflow.keras.utils.to_categorical', 'to_categorical', (['y_train'], {}), False, 'from tensorflow.keras.utils import to_categorical\n'), (45, 'tensorflow.keras.utils.to_categorical', 'to_categorical', (['y_valid'], {}), False, 'from tensorflow.keras.utils import to_categorical\n'), (20, 'os.path.join', 'os.path.join', (['dirpath', 'file'], {}), False, 'import os\n'), (21, 'cv2.imread', 'cv2.imread', (['filepath', 'cv2.IMREAD_UNCHANGED'], {}), False, 'import cv2\n'), (16, 'os.listdir', 'os.listdir', (['dirpath'], {}), False, 'import os\n'), (24, 'cv2.resize', 'cv2.resize', (['img', '(360, 363)', 'cv2.INTER_CUBIC'], {}), False, 'import cv2\n'), (37, 'numpy.array', 'np.array', (['X[:num_train]'], {}), True, 'import numpy as np\n'), (38, 'numpy.array', 'np.array', (['X[num_train:]'], {}), True, 'import numpy as np\n'), (39, 'numpy.array', 'np.array', (['y[:num_train]'], {}), True, 'import numpy as np\n'), (40, 'numpy.array', 'np.array', (['y[num_train:]'], {}), True, 'import numpy as np\n')]