# Copyright 2020 University of New South Wales, University of Sydney, Ingham Institute

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import SimpleITK as sitk


def detect_holes(img, lower_threshold=-10000, upper_threshold=-400):
    """
    Detect all (air) holes in given image. Default threshold values given are defined for CT.

    Args:
        img (sitk.Image): The image in which to detect holes.
        lower_threshold (int, optional): The lower threshold of intensities. Defaults to -10000.
        upper_threshold (int, optional): The upper threshold of intensities. Defaults to -400.

    Returns:
        label_image: image containing labels of all holes detected
        labels: a list of holes detected and some of their properties
    """

    # Threshold to get holes
    btif = sitk.BinaryThresholdImageFilter()
    btif.SetInsideValue(1)
    btif.SetOutsideValue(0)
    btif.SetLowerThreshold(lower_threshold)
    btif.SetUpperThreshold(upper_threshold)
    holes = btif.Execute(img)

    # Get connected components
    ccif = sitk.ConnectedComponentImageFilter()
    label_image = ccif.Execute(holes)

    labels = []

    lssif = sitk.LabelShapeStatisticsImageFilter()
    lssif.Execute(label_image)
    for region in range(1, ccif.GetObjectCount()):
        label = {
            "label": region,
            "phys_size": lssif.GetPhysicalSize(region),
            "elongation": lssif.GetElongation(region),
            "roundness": lssif.GetRoundness(region),
            "perimeter": lssif.GetPerimeter(region),
            "flatness": lssif.GetFlatness(region),
        }
        labels.append(label)

    # Sort by size
    labels = sorted(labels, key=lambda i: i["phys_size"], reverse=True)

    return label_image, labels


def get_external_mask(label_image, labels, kernel_radius=5):
    """Get a External mask given a label image and labels computed with detect_holes

    Args:
        label_image (sitk.Image): Label image from detect_holes
        labels (list): List of labels from detect_holes
        kernel_radius (int, optional): The width of the radius used by binary close. Defaults to 5.

    Returns:
        sitk.Image: External mask
    """

    bmcif = sitk.BinaryMorphologicalClosingImageFilter()
    bmcif.SetKernelType(sitk.sitkBall)
    bmcif.SetKernelRadius(kernel_radius)

    # Save the largest as the external contour
    external_mask = sitk.BinaryThreshold(label_image, labels[0]["label"], labels[0]["label"])
    external_mask = bmcif.Execute(external_mask)

    return external_mask


def get_lung_mask(label_image, labels, kernel_radius=2):
    """Get a Lung mask given a label image and labels computed with detect_holes

    Args:
        label_image (sitk.Image): Label image from detect_holes
        labels (list): List of labels from detect_holes
        kernel_radius (int, optional): The width of the radius used by binary close. Defaults to 2.

    Returns:
        sitk.Image: Combined lung mask
    """
    # Find potential lung regions (typically the 2nd and 3rd largest regions after the external air)
    lung_candidates = []
    for idx in range(1, min(5, len(labels))):  # Check among top 5 largest regions
        if 1.5 < labels[idx]["flatness"] < 3.0 and labels[idx]["phys_size"] > 1000:
            lung_candidates.append(labels[idx]["label"])
            if len(lung_candidates) == 2:  # Stop after finding 2 lung regions
                break

    if len(lung_candidates) < 2:
        print("Warning: Could not detect both lungs!")
        return None

    bmcif = sitk.BinaryMorphologicalClosingImageFilter()
    bmcif.SetKernelType(sitk.sitkBall)
    bmcif.SetKernelRadius(kernel_radius)

    # Create masks for both lungs and combine them
    lung_mask = sitk.BinaryThreshold(label_image, lung_candidates[0], lung_candidates[0])
    lung_mask2 = sitk.BinaryThreshold(label_image, lung_candidates[1], lung_candidates[1])
    lung_mask = sitk.Or(lung_mask, lung_mask2)
    lung_mask = bmcif.Execute(lung_mask)

    return lung_mask


def fill_holes(img, label_image, external_mask, lung_mask, fill_value=50):
    """Returns the input image with all holes filled (except for the external and lung holes).

    Args:
        img (sitk.Image): The image to fill
        label_image (sitk.Image): Label image from detect_holes
        external_mask (sitk.Image): The external mask of the patient.
        lung_mask (sitk.Image): The lung mask of the patient.
        fill_value (int): The value to use to fill the holes. Defaults to 50.

    Returns:
        sitk.Image: The image with holes filled
    """

    img_array = sitk.GetArrayFromImage(img)

    bdif = sitk.BinaryDilateImageFilter()
    bdif.SetKernelType(sitk.sitkBall)
    bdif.SetKernelRadius(3)

    mask = sitk.BinaryThreshold(label_image, 1, int(img_array.max()))
    mask = sitk.Subtract(mask, external_mask)
    mask = sitk.Subtract(mask, lung_mask)
    mask = bdif.Execute(mask)

    mask_array = sitk.GetArrayFromImage(mask)
    img_array[mask_array == 1] = fill_value

    filled_img = sitk.GetImageFromArray(img_array)
    filled_img.CopyInformation(img)

    return filled_img