import numpy as np
import matplotlib.pyplot as plt
import glob
import os
import sys
from PIL import Image
from tensorflow import keras
from keras.layers import Conv2D, Input, MaxPool2D,Flatten, Dense, Permute, GlobalAveragePooling2D
from keras.models import Model
# from keras.optimizers
from tensorflow.python.keras.optimizers import *
import numpy as np
import pickle
import keras
import cv2
import sys
import dlib
import os.path
import tensorflow as tf 
from tensorflow import keras
from tensorflow.keras.models import Sequential
from tensorflow.keras.applications.resnet50 import ResNet50
# from tensorflow.keras.applications.resnet50 import Dense
# from keras.optimizers import Adam
from tensorflow.keras.optimizers import Adam
import pickle
import numpy as np
import cv2
import os
from keras.layers import Dropout

masks = glob.glob("drive/MyDrive/saiannots1/*.png")
orgs = glob.glob("drive/MyDrive/saiframe1/*.jpg")


imgs_list = []
masks_list = []
for image, mask in zip(orgs, masks):
    imgs_list.append(np.array(Image.open(image).resize((256,256))))
    
    im = Image.open(mask).resize((256,256))
    masks_list.append(np.array(im))

imgs_np = np.asarray(imgs_list)
masks_np = np.asarray(masks_list)

x = np.asarray(imgs_np, dtype=np.float32)/255
y = np.asarray(masks_np, dtype=np.float32)/255


y = y.reshape(y.shape[0], y.shape[1], y.shape[2], 3)

print(x.shape, y.shape)

x = x.reshape(x.shape[0], x.shape[1], x.shape[2], 3)

print(x.shape, y.shape)


from sklearn.model_selection import train_test_split

x_train, x_val, y_train, y_val = train_test_split(x, y, test_size=0.5, random_state=0)
'''
print("x_train: ", x_train.shape)
print("y_train: ", y_train.shape)
print("x_val: ", x_val.shape)
print("y_val: ", y_val.shape)
'''

#Prepare train generator with data augmentation

# import get_augmented

train_gen = get_augmented(
    x_train, y_train, batch_size=2,
    data_gen_args = dict(
        rotation_range=15.,
        width_shift_range=0.05,
        height_shift_range=0.05,
        shear_range=50,
        zoom_range=0.2,
        horizontal_flip=True,
        vertical_flip=True,
        fill_mode='constant'
    ))

sample_batch = next(train_gen)
xx, yy = sample_batch
'''
print(xx.shape, yy.shape)
'''

'''
from keras_unet.utils import plot_imgs
plot_imgs(org_imgs=xx, mask_imgs=yy, nm_img_to_plot=2, figsize=6)
'''

#Initialize network

# from unet_model import unet_model

input_shape = x_train[0].shape

model = unet_model(
    input_shape,
    num_classes=1,
    filters=64,
    dropout=0.2,
    num_layers=4,
    output_activation='softmax'
)


print(model.summary())




#Compile + train

from keras.callbacks import ModelCheckpoint


model_filename = 'segm_model_v0.h5'
callback_checkpoint = ModelCheckpoint(
    model_filename, 
    verbose=1, 
    monitor='val_loss', 
    save_best_only=True,
)

from tensorflow.keras.optimizers import Adam, SGD
# from sklearn import metrics
from tensorflow.python.keras.metrics import Metric
# from tensorflow.keras.metrics import iou, iou_thresholded

model.compile(
    optimizer=Adam(learning_rate=0.001),
    loss='sparse_categorical_crossentropy',
    metrics=[iou, iou_thresholded,dice_coef,jaccard_coef]
    
)

history = model.fit(
    train_gen,
    steps_per_epoch=50,
    epochs=20,
    
    validation_data=(x_val, y_val),
    callbacks=[callback_checkpoint]
)