import os

import cv2

import numpy as np

import pandas as pd

from glob import glob

from tqdm import tqdm

import pydicom as dicom

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

""" Creating a directory """

def create_dir(path):

    if not os.path.exists(path):

        os.makedirs(path)

if __name__ == "__main__":

    """Seeding"""

    SEEDS = 42

    np.random.seed(SEEDS)

    tf.random.set_seed(SEEDS)

    """ Directory for storing files """

    create_dir("test")

    """ 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 = glob("data/test/*/*/*.dcm")

    print(f"Test: {len(test_x)}")

    """ Loop over the data """

    for x in tqdm(test_x):

        """Extract the Names"""

        dir_name = x.split("/")[-3]

        name = dir_name + "_" + x.split("/")[-1].split(".")[0]

        """ Read the Image """

        image = dicom.dcmread(x).pixel_array

        image = np.expand_dims(image, axis=-1)

        image = image / np.max(image) * 255.0

        x = image / 255.0

        x = np.concatenate([x, x, x], axis=-1)

        x = np.expand_dims(x, axis=0)

        """ Prediction """

        mask = model.predict(x)[0]

        mask = mask > 0.5

        mask = mask.astype(np.int32)

        mask = mask * 255

        cat_images = np.concatenate([image, mask], axis=1)

        cv2.imwrite(f"test/{name}.png", cat_images)