from collections import Counter

import pathlib

import math

import os

import matplotlib.pyplot as plt

import tensorflow as tf

import numpy as np

import cv2 as cv

import utils

import tqdm

def binary_image_from_dri(input_dir, threshold=128, save_dir=None):

    if os.path.isdir(input_dir):

        paths = utils.list_file(input_dir)

        utils.check_file([save_dir])

    else:

        paths = [input_dir]

    for path in paths:

        path_stem = pathlib.Path(path).stem

        image = cv.imread(path, 0)

        bin_image = binary_image(image, threshold)

        if save_dir is not None:

            cv.imwrite(os.path.join(save_dir, '{}.jpg'.format(path_stem)), bin_image)

    return

def binary_image(image, threshold):

    shape = image.shape

    if len(shape) == 3:

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

    th, bin_image = cv.threshold(image, threshold, 255, cv.THRESH_BINARY)

    return bin_image

def reverse_pred_image(normalize_pred_image):

    reverse_image = normalize_pred_image.squeeze() * 255.0

    reverse_image = np.array(reverse_image, dtype=np.uint8)

    return reverse_image

def save_images(pred, index, save_path, image_shape, split=False):

    image_numbers = int(np.sqrt(pred.shape[0]))

    if not split:

        h = image_shape[0]

        w = image_shape[1]

        H = int(image_numbers * image_shape[0])

        W = int(image_numbers * image_shape[1])

        big_image = np.zeros(shape=(H, W, image_shape[-1]), dtype=np.uint8).squeeze()

        for i in range(pow(image_numbers, 2)):

            image = (pred[i, :, :] * 255.0)

            image = np.array(image, dtype=np.uint8)

            image = image.squeeze()

            j = i % image_numbers

            k = i // image_numbers

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

                big_image[k * h:(k + 1) * h, j * w:(j + 1) * w] = image

            else:

                big_image[k * h:(k + 1) * h, j * w:(j + 1) * w, :] = image

        path = os.path.join(save_path, 'Segment_train_pred_{}.png'.format(index))

        plt.imsave(path, big_image, cmap='gray')

    else:

        for i in range(image_numbers ** 2):

            image = (pred[i, :, :] * 255.0)

            image = np.array(image, dtype=np.uint8)

            image = image.squeeze()

            path = os.path.join(save_path, 'Segment_pred_{}_{}.png'.format(index, i))

            plt.imsave(path, image, cmap='gray')

    return

def get_area(image):

    """Count the area of bleeding area"""

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

    _, bin_image = cv.threshold(image, 0, 255, cv.THRESH_BINARY)

    count_result = Counter(list(bin_image.reshape(-1, )))

    area = count_result.get(255)

    return area

def pixel_to_ml(pixel_area, dpi=96):

    if pixel_area is None:

        pixel_area = 0.0

    return pixel_area / pow(dpi, 2) * pow(25.4, 2) / 100

def draw_contours(image, mask, max_count=8, dpi=96):

    image = np.array(image)

    mask = np.array(mask)

    height, width = image.shape[:2]

    copy_image = image.copy()

    if len(mask.shape) == 3 and mask.shape[-1] != 1:

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

    else:

        mask = mask

    mask = cv.resize(mask, (height, width))

    th, bin_mask = cv.threshold(mask, 0, 255, cv.THRESH_BINARY)

    con_list = []

    blood_area = []

    contours, _ = cv.findContours(bin_mask, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)

    for index, contour in enumerate(contours):

        area = cv.contourArea(contour)

        if 200 < area < height * width * 0.94:

            con_list.append(index)

            blood_area.append(pixel_to_ml(area, dpi))

    if len(con_list) > max_count:

        blood_area = [0]

    else:

        for index in con_list:

            copy_image = cv.drawContours(copy_image, contours, index, (0, 0, 255), 5)

    return copy_image, sum(blood_area)

def save_invalid_data(origin_images, drawed_images, pred_mask_images, image_names, save_dir, reshape=True):

    """

     :param image_names: the image file names of the original images, in the form of a list

     :param origin_images: original bleeding images

     :param drawed_images: draw a contour map of the bleeding area on the original image according to the mask

     :param pred_mask_images: predicted mask image

     :param save_dir: save path of all images

     :param reshape: restore all images to the original image size

     """

    origin_save_dir = os.path.join(save_dir, 'origin')

    drawed_save_dir = os.path.join(save_dir, 'drawed')

    mask_save_dir = os.path.join(save_dir, 'pred_mask')

    utils.check_file([origin_save_dir, drawed_save_dir, mask_save_dir])

    for index in range(len(origin_images)):

        origin_image = origin_images[index]

        drawed_image = drawed_images[index]

        mask_image = pred_mask_images[index]

        _, bin_mask_image = cv.threshold(mask_image, 0, 255, cv.THRESH_BINARY)

        if reshape:

            origin_image = cv.resize(origin_image, (256, 256))

            drawed_image = cv.resize(drawed_image, (256, 256))

            save_name = '{}'.format(image_names[index])

            save_mask_path = os.path.join(mask_save_dir, save_name)

            save_origin_path = os.path.join(origin_save_dir, save_name)

            save_drawed_path = os.path.join(drawed_save_dir, save_name)

            cv.imwrite(save_mask_path, bin_mask_image)

            cv.imwrite(save_origin_path, origin_image)

            cv.imwrite(save_drawed_path, drawed_image)

    return

def count_volume(areas, thickness=0.45):

    for area in areas:

        if area == 0:

            areas.remove(area)

    areas_count = len(areas)

    volume = [areas[index] * thickness for index in range(areas_count)]

    return sum(volume)

def calculate_volume(mask_dir, real_shape=(1440, 1440), thickness=0.4, dpi=96):

    all_areas = []

    for path in tqdm.tqdm(pathlib.Path(mask_dir).iterdir()):

        mask_image = cv.imread(str(path))

        mask_image = cv.resize(mask_image, real_shape)

        area = pixel_to_ml(get_area(mask_image), dpi=dpi)

        all_areas.append(area)

    volume = count_volume(all_areas, thickness)

    return volume