import os
import cv2
from matplotlib import pyplot as plt
import numpy as np
import pickle
import gzip
import lzma
import pandas
import random
signnames = pandas.read_csv('signnames.csv')
%matplotlib inline
DATASET = 'train'
train_images = os.walk('dataset/'+DATASET)
train_image_list = []
PATH_TO_IMAGE = 'dataset/'+DATASET+'/'
RESIZE_IMG_SHAPE = 32
df = pandas.read_csv('germantraffic'+DATASET+'.csv')
for i in train_images:
    for item in i:
        if type(item) == str:
            SUBFOLDER_PATH = item
#             print('subfolder path',item)
        if type(item) == list:
            if len(item) > 50:
                for image in item:
                    if '.ppm' in image:
                        folder_name = SUBFOLDER_PATH+"/"+image
                        for index,filename in enumerate(df['filename']):
                            if filename == folder_name[-21:]:
                                for classindex, classname in enumerate(signnames['class']):
                                    if df['class'][index] == classname:
                                        label = classindex
                        current_image = cv2.imread(SUBFOLDER_PATH +"/" +  image)
                        current_image = current_image[:,:,::-1]
#                         current_image = current_image/255.0
                        h,w,a = current_image.shape
                        RESIZE_IMG_SHAPE = 32
                        blackbar_horizontal = 0
                        blackbar_vertical = 0
                        if h > w:
                            h_larger = True
                            resized_h = RESIZE_IMG_SHAPE
                            resized_w = int(RESIZE_IMG_SHAPE/h*w)
                            resized = cv2.resize(current_image,(resized_w,resized_h))
                            blackbar_horizontal = (RESIZE_IMG_SHAPE - resized_w)/2
                            for i in range(int(blackbar_horizontal)):
                                resized = np.insert(resized,[0,resized_w],0,axis=1)
                                resized_w+=1
                            if int(blackbar_horizontal / 0.5) % 2 == 1:
                                resized = np.insert(resized,0,0,axis=1)
                        if w > h:
                            w_larger = True
                            resized_w = RESIZE_IMG_SHAPE
                            resized_h = int(RESIZE_IMG_SHAPE/w*h)
                            resized = cv2.resize(current_image,(resized_w,resized_h))
                            blackbar_vertical = (RESIZE_IMG_SHAPE - resized_h)/2
                            for i in range(int(blackbar_vertical)):
                                resized = np.insert(resized,[0,resized_h],0,axis=0)
                                resized_h+=1
                            if int(blackbar_vertical / 0.5) % 2 == 1:
                                resized = np.insert(resized,0,0,axis=0)
                        if w == h:
                            w_equal_h = True
                            resized = cv2.resize(current_image,(RESIZE_IMG_SHAPE,RESIZE_IMG_SHAPE))
#                         train_image_list.append((resized.astype(np.float16), label))
                        train_image_list.append((resized,label))
#                         train_image_list.append(resized.astype(np.float16))

print(len(train_image_list))
random.shuffle(train_image_list)
pickle_out = lzma.open(f"{DATASET}_data_not_normalized.pickle","wb",preset=9)
pickle.dump(train_image_list, pickle_out)
pickle_out.close()
print("Done!")