import os
import cv2
import numpy as np
from glob import glob
from sklearn.utils import shuffle

import tensorflow as tf
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.metrics import Recall, Precision

from models.unet import get_unet_model
from metrics import dice_loss, dice_coef, iou

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

IMG_HEIGHT = 512
IMG_WIDTH = 512
AUTO = tf.data.AUTOTUNE


def create_dir(path):
    """Create a directory."""
    if not os.path.exists(path):
        os.makedirs(path)


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


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


def preprocess_image(path):
    path = path.decode()
    x = cv2.imread(path, cv2.IMREAD_COLOR)
    x = x / 255.0
    x = x.astype(np.float32)
    return x


def preprocess_mask(path):
    path = path.decode()
    x = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
    x = x / 255.0
    x = x > 0.5
    x = x.astype(np.float32)
    x = np.expand_dims(x, axis=-1)
    return x


def preprocess_data(x, y):
    def _parse(x, y):
        x = preprocess_image(x)
        y = preprocess_mask(y)
        return x, y

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


def load_dataset(x, y, batch_size=8):
    dataset = tf.data.Dataset.from_tensor_slices((x, y))
    dataset = (
        dataset.map(preprocess_data, num_parallel_calls=AUTO)
        .batch(batch_size)
        .prefetch(AUTO)
    )
    return dataset


if __name__ == "__main__":
    """Seeding"""
    SEEDS = 42
    np.random.seed(SEEDS)
    tf.random.set_seed(SEEDS)

    # Create a MirroredStrategy.
    strategy = tf.distribute.MirroredStrategy()
    print("Number of devices: {}".format(strategy.num_replicas_in_sync))

    """ Directory for storing files """
    create_dir("files")

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

    """ Dataset """
    dataset_path = os.path.join("new_data")
    train_path = os.path.join(dataset_path, "train")
    valid_path = os.path.join(dataset_path, "valid")

    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 = load_dataset(train_x, train_y, batch_size=batch_size)
    valid_dataset = load_dataset(valid_x, valid_y, batch_size=batch_size)

    """ Model """
    # Open a strategy scope.
    with strategy.scope():
        model = get_unet_model((IMG_HEIGHT, IMG_WIDTH, 3))
        metrics = [dice_coef, iou, Recall(), Precision()]
        model.compile(loss=dice_loss, optimizer=Adam(lr), metrics=metrics)

    """Setting up Training Callbacks"""
    train_callbacks = [
        tf.keras.callbacks.ModelCheckpoint(model_path, verbose=1, save_best_only=True),
        tf.keras.callbacks.ReduceLROnPlateau(
            monitor="val_loss", factor=0.1, patience=10, min_lr=1e-7, verbose=1
        ),
        tf.keras.callbacks.CSVLogger(csv_path),
        tf.keras.callbacks.TensorBoard(),
        tf.keras.callbacks.EarlyStopping(
            monitor="val_loss", patience=10, restore_best_weights=False
        ),
    ]

    history = model.fit(
        train_dataset,
        epochs=num_epochs,
        validation_data=valid_dataset,
        callbacks=train_callbacks,
        shuffle=False,
    )