import numpy as np
import torch

def brats_segmentation_regions() -> dict:
    return {"ET_brats": 4, "ET": 3,  "NCR-NET": 1, "ED": 2}


def get_ncr_net(segmentation_map: np.ndarray) -> np.ndarray:
    regions = brats_segmentation_regions()
    copied_segmentation = _copy_input(segmentation_map)
    copied_segmentation[copied_segmentation != regions["NCR-NET"]] = 0
    return copied_segmentation.astype(np.uint8)


def get_ed(segmentation_map: np.ndarray) -> np.ndarray:
    regions = brats_segmentation_regions()
    copied_segmentation = _copy_input(segmentation_map)
    copied_segmentation[copied_segmentation != regions["ED"]] = 0
    return copied_segmentation.astype(np.uint8)


def get_et(segmentation_map: np.ndarray) -> np.ndarray:
    """
    ET : enhancing tumors is label 3 in the code, 4 as brats
    :param segmentation_map:
    :return: only label for ET
    """
    regions = brats_segmentation_regions()
    copied_segmentation = _copy_input(segmentation_map)
    unique_values = np.unique(copied_segmentation)
    if max(unique_values) == 3:
        copied_segmentation[copied_segmentation != regions["ET"]] = 0
    else:
        copied_segmentation[copied_segmentation != regions["ET_brats"]] = 0

    copied_segmentation[copied_segmentation != 0] = 1
    return copied_segmentation.astype(np.uint8)


def get_wt(segmentation_map: np.ndarray) -> np.ndarray:
    """ WT : entails all regions 4 (ET, NCR, ED ) """
    copied_segmentation = _copy_input(segmentation_map)
    copied_segmentation[copied_segmentation != 0] = 1
    return copied_segmentation.astype(np.uint8)


def get_tc(segmentation_map: np.ndarray) -> np.ndarray:
    """ TC : tumor core entails the ET and NCR/NET labels 4 and 1 """
    regions = brats_segmentation_regions()

    copied_segmentation = _copy_input(segmentation_map)
    copied_segmentation[copied_segmentation == regions["ED"]] = 0  # remove edema
    copied_segmentation[copied_segmentation > 0] = 1

    return copied_segmentation.astype(np.uint8)

def convert_from_brats_labels(segmentation_map: np.ndarray) -> np.ndarray:
    """Method to convert brats labels as models need consecutive values"""
    regions = brats_segmentation_regions()
    segmentation_map[segmentation_map == regions["ET_brats"]] = regions["ET"]
    return segmentation_map

def convert_to_brats_labels(segmentation_map: np.ndarray) -> np.ndarray:
    """Method to convert recover brats labels encoding"""
    regions = brats_segmentation_regions()
    segmentation_map[segmentation_map == regions["ET"]] = regions["ET_brats"]
    return segmentation_map


def _copy_input(input):
    if torch.is_tensor(input):
        return input.detach().clone()
    else:
        return input.copy()
