import tensorflow as tf
from sklearn.model_selection import train_test_split
from PIL import Image
import os
import numpy as np
import random
import cv2
import model
import traitement_images as ti

dir_images='./training/images/'
dir_mask  ='./training/1st_manual/'

if not os.path.isdir(dir_images):
    quit("The directory {} doesn't exist !".format(dir_images))
if not os.path.isdir(dir_mask):
    quit("The directory {} doesn't exist !".format(dir_mask))

tab_images=[]
tab_masks=[]

list_file=os.listdir(dir_images)
if list_file is None:
    quit("No file in {} !".format(dir_images))
    
for fichier in list_file:
    img_orig=cv2.imread(dir_images+fichier)
    tab_images.append(img_orig[:576, :560])
    num=fichier.split('_')[0]
    file_mask=dir_mask+num+'_manual1.gif'
    if not os.path.isfile(file_mask):
        quit("Mask of {} doesn't exist in {}".format(file_mask, dir_mask))
    img_mask_orig=np.array(Image.open(file_mask))
    tab_masks.append(img_mask_orig[:576, :560])

    for angle in range(0, 360, 30):
        img_r=ti.rotateImage(img_orig, angle)
        img=img_r.copy()
        img=ti.random_change(img)
        tab_images.append(img[:576, :560])
        img_mask=ti.rotateImage(img_mask_orig, angle)
        tab_masks.append(img_mask[:576, :560])
        
        img=cv2.flip(img_r, 0)
        img=ti.random_change(img)
        tab_images.append(img[:576, :560])
        img_m=cv2.flip(img_mask, 0)
        tab_masks.append(img_m[:576, :560])

        img=cv2.flip(img_r, 1)
        img=ti.random_change(img)
        tab_images.append(img[:576, :560])
        img_m=cv2.flip(img_mask, 1)
        tab_masks.append(img_m[:576, :560])

        img=cv2.flip(img_r, -1)
        img=ti.random_change(img)
        tab_images.append(img[:576, :560])
        img_m=cv2.flip(img_mask, -1)
        tab_masks.append(img_m[:576, :560])

tab_images=np.array(tab_images, dtype=np.float32)/255
tab_masks =np.array(tab_masks,  dtype=np.float32)[:, :, :]/255

train_images, test_images, train_masks, test_masks=train_test_split(tab_images, tab_masks, test_size=0.05)

del tab_images
del tab_masks

my_model=model.model(64)

my_model.compile(optimizer='adam',
                 loss='binary_crossentropy',
                 metrics=['accuracy'])
my_model.fit(train_images,
             train_masks,
             epochs=20,
             batch_size=4,
             validation_data=(test_images, test_masks))

dir_test_images='./test/images/'

tab_test_images=[]
tab_files=[]
for fichier in os.listdir(dir_test_images):
    img=cv2.imread(dir_test_images+fichier)
    tab_test_images.append(img[:576, :560])
    tab_files.append(fichier.split('_')[0])

tab_test_images=np.array(tab_test_images, dtype=np.float32)/255
tab_files=np.array(tab_files)

for id in range(len(tab_test_images)):
    mask=np.zeros((584, 565, 1), dtype=np.float32)
    prediction=my_model.predict(np.array([tab_test_images[id]]))
    mask[:576, :560]=prediction[0]*255
    cv2.imwrite("./predictions/"+str(tab_files[id])+".png", mask)