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)