import os
import cv2
import numpy as np
import pandas as pd
from glob import glob
from tqdm import tqdm

import tensorflow as tf
from tensorflow.keras.utils import CustomObjectScope
from sklearn.metrics import (
    accuracy_score,
    f1_score,
    jaccard_score,
    precision_score,
    recall_score,
)
from metrics import dice_loss, dice_coef, iou
from train import load_data

IMG_HEIGHT = 512
IMG_WIDTH = 512

""" Creating a directory """


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


def save_results(image, mask, y_pred, save_image_path):
    ## i - m - y
    line = np.ones((IMG_HEIGHT, 10, 3)) * 128

    """ Mask """
    mask = np.expand_dims(mask, axis=-1)  ## (512, 512, 1)
    mask = np.concatenate([mask, mask, mask], axis=-1)  ## (512, 512, 3)

    """ Predicted Mask """
    y_pred = np.expand_dims(y_pred, axis=-1)  ## (512, 512, 1)
    y_pred = np.concatenate([y_pred, y_pred, y_pred], axis=-1)  ## (512, 512, 3)
    y_pred = y_pred * 255

    cat_images = np.concatenate([image, line, mask, line, y_pred], axis=1)
    cv2.imwrite(save_image_path, cat_images)


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

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

    """ Loading model """
    with CustomObjectScope(
        {"iou": iou, "dice_coef": dice_coef, "dice_loss": dice_loss}
    ):
        model = tf.keras.models.load_model("files/model.h5")

    """ Load the dataset """
    test_x = sorted(glob(os.path.join("new_data", "valid", "image", "*")))
    test_y = sorted(glob(os.path.join("new_data", "valid", "mask", "*")))
    print(f"Test: {len(test_x)} - {len(test_y)}")

    """ Evaluation and Prediction """
    SCORE = []
    for x, y in tqdm(zip(test_x, test_y), total=len(test_x)):
        """Extract the name"""
        name = x.split("/")[-1].split(".")[0]

        """ Reading the image """
        image = cv2.imread(x, cv2.IMREAD_COLOR)
        x = image / 255.0
        x = np.expand_dims(x, axis=0)

        """ Reading the mask """
        mask = cv2.imread(y, cv2.IMREAD_GRAYSCALE)
        y = mask / 255.0
        y = y > 0.5
        y = y.astype(np.int32)

        """ Prediction """
        y_pred = model.predict(x)[0]
        y_pred = np.squeeze(y_pred, axis=-1)
        y_pred = y_pred > 0.5
        y_pred = y_pred.astype(np.int32)

        """ Saving the prediction """
        save_image_path = f"results/{name}.png"
        save_results(image, mask, y_pred, save_image_path)

        """ Flatten the array """
        y = y.flatten()
        y_pred = y_pred.flatten()

        """ Calculating the metrics values """
        acc_value = accuracy_score(y, y_pred)
        f1_value = f1_score(y, y_pred, labels=[0, 1], average="binary", zero_division=1)
        jac_value = jaccard_score(
            y, y_pred, labels=[0, 1], average="binary", zero_division=1
        )
        recall_value = recall_score(
            y, y_pred, labels=[0, 1], average="binary", zero_division=1
        )
        precision_value = precision_score(
            y, y_pred, labels=[0, 1], average="binary", zero_division=1
        )
        SCORE.append(
            [name, acc_value, f1_value, jac_value, recall_value, precision_value]
        )

    """ Metrics values """
    score = [s[1:] for s in SCORE]
    score = np.mean(score, axis=0)
    print(f"Accuracy: {score[0]:0.5f}")
    print(f"F1: {score[1]:0.5f}")
    print(f"Jaccard: {score[2]:0.5f}")
    print(f"Recall: {score[3]:0.5f}")
    print(f"Precision: {score[4]:0.5f}")

    df = pd.DataFrame(
        SCORE, columns=["Image", "Accuracy", "F1", "Jaccard", "Recall", "Precision"]
    )
    df.to_csv("files/score.csv")