import numpy as np
import SimpleITK as sitk

channels = [
    "background",
    "spleen",
    "right_kidney",
    "left_kidney",
    "gallbladder",
    "liver",
    "stomach",
    "pancreas",
    "right_adrenal_gland",
    "left_adrenal_gland",
    "left_lung",
    "right_lung",
    "heart",
    "aorta",
    "inferior_vena_cava",
    "portal_vein_and_splenic_vein",
    "left_iliac_artery",
    "right_iliac_artery",
    "left_iliac_vena",
    "right_iliac_vena",
    "esophagus",
    "small_bowel",
    "duodenum",
    "colon",
    "urinary_bladder",
    "spine",
    "sacrum",
    "left_hip",
    "right_hip",
    "left_femur",
    "right_femur",
    "left_autochthonous_muscle",
    "right_autochthonous_muscle",
    "left_iliopsoas_muscle",
    "right_iliopsoas_muscle",
    "left_gluteus_maximus",
    "right_gluteus_maximus",
    "left_gluteus_medius",
    "right_gluteus_medius",
    "left_gluteus_minimus",
    "right_gluteus_minimus",
]


def make_isotropic(image, interpolator=sitk.sitkLinear, spacing=None):
    """
    Many file formats (e.g. jpg, png,...) expect the pixels to be isotropic, same
    spacing for all axes. Saving non-isotropic data in these formats will result in
    distorted images. This function makes an image isotropic via resampling, if needed.
    Args:
        image (SimpleITK.Image): Input image.
        interpolator: By default the function uses a linear interpolator. For
                      label images one should use the sitkNearestNeighbor interpolator
                      so as not to introduce non-existant labels.
        spacing (float): Desired spacing. If none given then use the smallest spacing from
                         the original image.
    Returns:
        SimpleITK.Image with isotropic spacing which occupies the same region in space as
        the input image.
    """
    original_spacing = image.GetSpacing()
    # Image is already isotropic, just return a copy.
    if all(spc == original_spacing[0] for spc in original_spacing):
        return sitk.Image(image)
    # Make image isotropic via resampling.
    original_size = image.GetSize()
    if spacing is None:
        spacing = min(original_spacing)
    new_spacing = [spacing] * image.GetDimension()
    new_size = [int(round(osz * ospc / spacing)) for osz, ospc in zip(original_size, original_spacing)]
    return sitk.Resample(
        image,
        new_size,
        sitk.Transform(),
        interpolator,
        image.GetOrigin(),
        new_spacing,
        image.GetDirection(),
        0,  # default pixel value
        image.GetPixelID(),
    )


def label_mapper(seg):

    labels = []
    for _class in np.unique(seg):
        if _class == 0:
            continue
        labels.append((seg == _class, channels[_class]))

    return labels


def sitk2numpy(img, normalize=False):
    img = sitk.DICOMOrient(img, "LPS")
    # img = make_isotropic(img)
    img = sitk.GetArrayFromImage(img)
    if normalize:
        minval, maxval = np.min(img), np.max(img)
        img = ((img - minval) / (maxval - minval)).clip(0, 1) * 255
    img = img.astype(np.uint8)
    return img


def read_image(path, normalize=False):

    img = sitk.ReadImage(path)
    return sitk2numpy(img, normalize)


def display(image, seg=None, _slice=50):

    # Image
    if image is None or (isinstance(image, list) and len(image) == 0):
        return None
    if isinstance(image, list):
        image = image[-1]
    x = int(_slice * (image.shape[0] / 100))
    image = image[x, :, :]

    # Segmentation
    if seg is None or (isinstance(seg, list) and len(seg) == 0):
        seg = []
    else:
        if isinstance(seg, list):
            seg = seg[-1]
        seg = label_mapper(seg[x, :, :])

    return image, seg


def read_and_display(path, image_state, seg_state):

    image_state.clear()
    seg_state.clear()

    if path is not None:
        image = read_image(path, normalize=True)
        image_state.append(image)
        return display(image), image_state, seg_state
    else:
        return None, image_state, seg_state