import os
import tensorflow as tf
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
import numpy as np
import cv2 
from glob import glob 
from sklearn.utils import shuffle
from tensorflow.keras.callbacks import ModelCheckpoint,CSVLogger,ReduceLROnPlateau
from tensorflow.keras.optimizers import Adam 
from tensorflow.keras.metrics import Recall,Precision
from model_copy import build_unet 
from metrics import dice_loss,dice_coef,iou 


H = 512 
W = 512 

def create_dir(path):
    if not os.path.exists(path):
        os.makedirs(path)

def load_data(path):
    x = sorted(glob(os.path.join(path,"images","*.jpg")))
    y = sorted(glob(os.path.join(path,"mask","*.jpg")))
    
    return x,y

def shuffling(x,y):
    x,y = shuffle(x,y,random_state = 42)
    
    return x,y

def read_image(path):
    path = path.decode()
    x = cv2.imread(path,cv2.IMREAD_COLOR)
    # x = cv2.resize(x,(W,H))
    x = x/255.0
    x = x.astype(np.float32)
    return x 

def read_mask(path):
    path = path.decode()
    x = cv2.imread(path,cv2.IMREAD_GRAYSCALE)
    # x = cv2.resize(x,(W,H))
    x = x/255.0
    x = x.astype(np.float32)
    x = np.expand_dims(x,axis=-1)
    return x 

def tf_parse(x,y):
    def _parse(x,y):
        x = read_image(x)
        y = read_mask(y)
        return x,y

    x,y = tf.numpy_function(_parse,[x,y],[tf.float32,tf.float32])
    x.set_shape([H,W,3])
    y.set_shape([H,W,1])
    return x,y 

def tf_dataset(X,Y,batch_size=2):
    dataset = tf.data.Dataset.from_tensor_slices((X,Y))
    dataset = dataset.map(tf_parse)
    dataset = dataset.batch(batch_size)
    dataset = dataset.prefetch(4)
    return dataset












if __name__ == "__main__":


    '''Seeding'''

    np.random.seed(42)
    tf.random.set_seed(42)

    ''' Directory to save files '''

    create_dir("files")

    """Hyperparameters"""

    batch_size = 2
    lr = 1e-4
    num_epochs = 100
    model_path = os.path.join("files","model.h5")
    csv_path = os.path.join("files","data.csv")

    """ Dataset """
    dataset_path = "newdata"
    train_path = os.path.join(dataset_path,"train")
    valid_path = os.path.join(dataset_path,"test")
    

    train_x,train_y = load_data(train_path)
    
    train_x,train_y = shuffling(train_x,train_y)
    
    valid_x,valid_y = load_data(valid_path)
    
    print(f"Train:{len(train_x)}-{len(train_y)}")
    print(f"Valid:{len(valid_x)}-{len(valid_y)}")


    train_dataset = tf_dataset(train_x,train_y,batch_size=batch_size)
    valid_dataset = tf_dataset(valid_x,valid_y,batch_size=batch_size)

    train_steps = len(train_x)//batch_size
    valid_steps = len(valid_x)//batch_size

    if len(train_x) % batch_size != 0 :
        train_steps += 1
    if len(train_x) % batch_size != 0 :
        train_steps += 1
    
    """ Model """ 
    model = build_unet((H,W,3))
    model.compile(loss=dice_loss,optimizer=Adam(lr),metrics=[dice_coef,iou,Recall])
    model.summary()

    callbacks = [
        ModelCheckpoint(model_path,verbose=1,save_best_only=True),
        ReduceLROnPlateau(monitor="val_loss",factor=0.1,patience=5,min_lr= 1e-6,verbose=1),
        CSVLogger(csv_path),
        TensorBoard(),
        EarlyStopping(monitor = "val_loss",patience=10,restore_best_weights=False)


    ]

    model.fit(train_dataset,epochs=num_epochs,validation_data=valid_dataset,steps_per_epochs=train_steps,validation_steps=valid_steps,callbacks=callbacks)






