import matplotlib.pyplot as plt
import numpy as np
import csv
import os

class DataLoder:

    def load_data(self, train_path, test_path, label_path):
        # Please make sure to change this function to load your train/validation/test data.
        # train_data = np.array([plt.imread('./data/test_images/0.jpg'), plt.imread('./data/test_images/1.jpg'),
        #               plt.imread('./data/test_images/2.jpg'), plt.imread('./data/test_images/3.jpg')])
    
        trainPath = train_path
        #trainPath = "./100_3dimage/"
        
        trainDir = os.listdir(trainPath)
        train_data = np.array(plt.imread(trainPath + trainDir[0])).reshape(1, 224, 224, 3)
        classData = {}
        
        with open(label_path, mode='r', encoding="UTF8") as f :
            data = csv.reader(f)
            for row in data :
                classData[row[0] + ".jpg"] = row[1]
        
        # with open('./trainSet.csv', mode='r', encoding="UTF8") as f :
        #     data = csv.reader(f)
        #     for row in data :
        #         classData[row[0].replace(".vir", ".jpg").replace('\ufeff','')] = row[1]

        classLabel = [int(classData[trainDir[0]])]
        i = 0
        for fil in trainDir[1:] :
            i += 1
            tempArry = np.array(plt.imread(trainPath + fil)).reshape(1, 224, 224, 3)
            train_data = np.append(train_data, tempArry, axis=0)
            classLabel.append(int(classData[fil]))
            if i % 100 == 0 :
                print(i, "files now loaded!")
            

        print(train_data.shape)

        # 인풋 데이터
        X_train = train_data
        # 트레인 데이터 라벨
        y_train = np.array(classLabel)
        print("train label")
        
        valPath = test_path

        valDir = os.listdir(valPath)
        val_data = np.array(plt.imread(valPath + valDir[0])).reshape(1, 224, 224, 3)

        valClassLabel = [int(classData[valDir[0]])]
        i = 0
        for fil in valDir[1:] :
            i += 1
            tempArry = np.array(plt.imread(valPath + fil)).reshape(1, 224, 224, 3)
            val_data = np.append(val_data, tempArry, axis=0)
            valClassLabel.append(int(classData[fil]))
            if i % 100 == 0 :
                print(i, "files now loaded!")
            

        print(val_data.shape)

        X_val = val_data
        y_val = np.array(valClassLabel)
        
        return X_train.astype('float32'), X_val.astype('float32'), y_train,  y_val