import os

import pathlib

import tqdm

import cv2 as cv

import numpy as np

import tensorflow as tf

from tensorflow.keras import *

import matplotlib.pyplot as plt

import tensorflow.keras as keras

import utils

def return_inputs(inputs):

    """Returns the output value according to the input type, used for image path input"""

    all_image_paths = None

    if type(inputs) is str:

        if os.path.isfile(inputs):

            all_image_paths = [inputs]

        elif os.path.isdir(inputs):

            all_image_paths = utils.list_file(inputs)

    elif type(inputs) is list:

        all_image_paths = inputs

    return all_image_paths

# 1. make dataset

def get_path_name(data_dir, get_id=False, nums=-1):

    name_list = []

    path_list = []

    for path in pathlib.Path(data_dir).iterdir():

        path_list.append(str(path))

        if get_id:

            name_list.append(path.stem[-5:])

        else:

            name_list.append(path.stem)

    if nums != -1:

        name_list = name_list[:nums]

        path_list = path_list[:nums]

    name_list = sorted(name_list, key=lambda path_: int(pathlib.Path(path_).stem))

    path_list = sorted(path_list, key=lambda path_: int(pathlib.Path(path_).stem))

    return name_list, path_list

class TFData:

    def __init__(self, image_shape, image_dir=None, mask_dir=None,

                 out_name=None, out_dir='', zip_file=True, mask_gray=True):

        self.image_shape = image_shape

        self.zip_file = zip_file

        self.image_dir = image_dir

        self.mask_dir = mask_dir

        self.out_name = out_name

        self.out_dir = os.path.join(out_dir, out_name)

        self.mask_gray = mask_gray

        if len(image_shape) == 3 and image_shape[-1] != 1:

            self.image_gray = False

        else:

            self.image_gray = True

        if self.zip_file:

            self.options = tf.io.TFRecordOptions(compression_type='GZIP')

        if image_dir is not None and mask_dir is not None:

            self.image_name, self.image_list = get_path_name(self.image_dir, False)

            self.mask_name, self.mask_list = get_path_name(self.mask_dir, False)

            self.data_zip = zip(self.image_list, self.mask_list)

    def image_to_byte(self, path, gray_scale):

        image = cv.imread(path)

        if not gray_scale:

            image = cv.cvtColor(image, cv.COLOR_BGR2RGB)

        elif len(image.shape) == 3:

            image = cv.cvtColor(image, cv.COLOR_BGR2GRAY)

        else:

            pass

        image = cv.resize(image, tuple(self.image_shape[:2]))

        return image.tobytes()

    def write_tfrecord(self):

        if not os.path.exists(self.out_dir):

            if self.zip_file:

                writer = tf.io.TFRecordWriter(self.out_dir, self.options)

            else:

                writer = tf.io.TFRecordWriter(self.out_dir)

            print(len(self.image_list))

            for image_path, mask_path in tqdm.tqdm(self.data_zip, total=len(self.image_list)):

                image = self.image_to_byte(image_path, self.image_gray)

                mask = self.image_to_byte(mask_path, self.mask_gray)

                example = tf.train.Example(features=tf.train.Features(

                    feature={

                        'mask': tf.train.Feature(bytes_list=tf.train.BytesList(value=[mask])),

                        'image': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image]))

                    }

                ))

                writer.write(example.SerializeToString())

            writer.close()

        print('Dataset finished!')

    def _parse_function(self, example_proto):

        features = tf.io.parse_single_example(

            example_proto,

            features={

                'mask': tf.io.FixedLenFeature([], tf.string),

                'image': tf.io.FixedLenFeature([], tf.string)

            }

        )

        image = features['image']

        image = tf.io.decode_raw(image, tf.uint8)

        if self.image_gray:

            image = tf.reshape(image, self.image_shape[:2])

            image = tf.expand_dims(image, -1)

        else:

            image = tf.reshape(image, self.image_shape)

        label = features['mask']

        label = tf.io.decode_raw(label, tf.uint8)

        if self.mask_gray:

            label = tf.reshape(label, self.image_shape[:2])

            label = tf.expand_dims(label, -1)

        else:

            label = tf.reshape(label, self.image_shape)

        return image, label

    def data_iterator(self, batch_size, data_name='', repeat=1, shuffle=True):

        if len(data_name) == 0:

            data_name = self.out_dir

        else:

            data_name = data_name

        if self.zip_file:

            dataset = tf.data.TFRecordDataset(data_name, compression_type='GZIP')

        else:

            dataset = tf.data.TFRecordDataset(data_name)

        dataset = dataset.map(self._parse_function)

        if shuffle:

            dataset = dataset.shuffle(buffer_size=100).repeat(repeat).batch(batch_size, drop_remainder=True)

        else:

            dataset = dataset.repeat(repeat).batch(batch_size, drop_remainder=True)

        return dataset

def data_preprocess(image, mask):

    """Normalize the image and mask data sets between 0-1"""

    image = tf.cast(image, np.float32)

    image = image / 127.5 - 1

    mask = tf.cast(mask, np.float32)

    mask = mask / 255.0

    return image, mask

def make_data(image_shape, image_dir, mask_dir, out_name=None, out_dir=''):

    tf_data = TFData(image_shape=image_shape, out_dir=out_dir, out_name=out_name,

                     image_dir=image_dir, mask_dir=mask_dir)

    tf_data.write_tfrecord()

    return

def get_tfrecord_data(tf_record_path, tf_record_name, data_shape, batch_size=32, repeat=1, shuffle=True):

    tf_data = TFData(image_shape=data_shape, out_dir=tf_record_path, out_name=tf_record_name)

    seg_data = tf_data.data_iterator(batch_size=batch_size, repeat=repeat, shuffle=shuffle)

    seg_data = seg_data.map(data_preprocess)

    return seg_data

def get_test_data(test_data_path, image_shape, image_nums=16):

    """

    :param test_data_path: test image path

    :param image_shape: Need to resize the shape of the test image, a tuple of length 3, [height, width, channel]

    :param image_nums: How many images need to be tested, the default is 16

    :return: normalized image collection

    """

    or_resize_shape = (1440, 1440)

    normalize_test_data = []

    original_test_data = []

    test_image_name = []

    test_data_paths = return_inputs(test_data_path)

    for path in test_data_paths:

        try:

            test_image_name.append(pathlib.Path(path).name)

            original_test_image = cv.imread(str(path))

            original_test_image = cv.resize(original_test_image, or_resize_shape)

            original_shape = original_test_image.shape

            if len(original_shape) == 0:

                print('Unable to read the {} file, please keep the path without Chinese! --First'.format(str(path)))

            else:

                original_test_data.append(original_test_image)

            if image_shape[-1] == 1:

                original_test_image = cv.cvtColor(original_test_image, cv.COLOR_BGR2GRAY)

            image = cv.resize(original_test_image, tuple(image_shape[:2]))

            image = image.astype(np.float32)

            image = image / 127.5 - 1

            normalize_test_data.append(image)

            if image_nums == -1:

                pass

            else:

                if len(normalize_test_data) == image_nums:

                    break

        except Exception as e:

            print('Unable to read the {} file, please keep the path without Chinese! --Second'.format(str(path)))

            print(e)

    normalize_test_array = np.array(normalize_test_data)

    if image_shape[-1] == 1:

        normalize_test_array = np.expand_dims(normalize_test_array, -1)

    original_test_array = np.array(original_test_data)

    if original_test_array.shape == 3:

        original_test_array = np.expand_dims(original_test_array, 0)

        normalize_test_array = np.expand_dims(normalize_test_array, 0)

    return test_image_name, original_test_array, normalize_test_array