'''

Copyright (c) Microsoft Corporation. All rights reserved.

Licensed under the MIT License.

'''

import SimpleITK as sitk 

import numpy as np  

import cc3d

#%%

def get_3darray_from_niftipath(

    path: str,

) -> np.ndarray:

    """Get a numpy array of a Nifti image using the filepath

    Args:

        path (str): path of the Nifti file

    Returns:

        np.ndarray: 3D numpy array for the image

    """

    image = sitk.ReadImage(path)

    array = np.transpose(sitk.GetArrayFromImage(image), (2,1,0))

    return array

def calculate_patient_level_lesion_suvmean_suvmax(

    ptarray: np.ndarray, 

    maskarray: np.ndarray,

    marker: str = 'SUVmean'

) -> np.float64:

    """Function to return the lesion SUVmean or SUVmax for all lesions in 

    a 3D PET image using the corresponding 3D segmentation mask 

    Args:

        ptarray (np.ndarray): numpy ndarray for 3D PET image

        maskarray (np.ndarray): numpy ndarray for 3D mask image

        marker (str, optional): Whether you want to calculate SUVmean or SUVmax . 

        Defaults to 'SUVmean'.

    Returns:

        np.float64: patient-level SUVmean or SUVmax

    """

    prod = np.multiply(ptarray, maskarray)

    num_nonzero_voxels = len(np.nonzero(maskarray)[0])

    if num_nonzero_voxels == 0:

        return 0.0

    else:

        if marker == 'SUVmean':

            return np.sum(prod)/num_nonzero_voxels

        elif marker == 'SUVmax':

            return np.max(prod)

#%%

def calculate_patient_level_tmtv(

    maskarray: np.ndarray,

    spacing: tuple

) -> np.float64:

    """Function to return the total metabolic tumor volume (TMTV) in cm^3 using 

    3D mask containing 0s for background and 1s for lesions/tumors

    Args:

        maskarray (np.ndarray): numpy ndarray for 3D mask image

    Returns:

        np.float64: 

    """

    voxel_volume_cc = np.prod(spacing)/1000 # voxel volume in cm^3

    num_lesion_voxels = len(np.nonzero(maskarray)[0])

    tmtv_cc = voxel_volume_cc*num_lesion_voxels

    return tmtv_cc

#%%

def calculate_patient_level_lesion_count(

    maskarray: np.ndarray,

) -> int:

    """Function to return the total number of lesions using the 3D segmentation mask 

    Args:

        maskarray (np.ndarray): numpy ndarray for 3D mask image

    Returns:

        int: _description_

    """

    _, num_lesions = cc3d.connected_components(maskarray, connectivity=18, return_N=True)

    return num_lesions

#%%

def calculate_patient_level_tlg(

    ptarray: np.ndarray,

    maskarray: np.ndarray,

    spacing: tuple

) -> np.float64:

    """Function to return the total lesion glycolysis (TLG) using a 3D PET image 

    and the corresponding 3D segmentation mask (containing 0s for background and

    1s for lesion/tumor)

    TLG = SUV1*V1 + SUV2*V2 + ... + SUVn*Vn, where SUV1...SUVn are the SUVmean 

    values of lesions 1...n with volumes V1...Vn, respectively

    Args:

        ptarray (np.ndarray): numpy ndarray for 3D PET image

        maskarray (np.ndarray): numpy ndarray for 3D mask image

    Returns:

        np.float64: total lesion glycolysis in cm^3 (assuming SUV is unitless)

    """

    voxel_volume_cc = np.prod(spacing)/1000 # voxel volume in cm^3

    labels_out, num_lesions = cc3d.connected_components(maskarray, connectivity=18, return_N=True)

    if num_lesions == 0:

        return 0.0

    else:

        _, lesion_num_voxels = np.unique(labels_out, return_counts=True)

        lesion_num_voxels = lesion_num_voxels[1:]

        lesion_mtvs = voxel_volume_cc*lesion_num_voxels

        lesion_suvmeans = []

        for i in range(1, num_lesions+1):

            mask = np.zeros_like(labels_out)

            mask[labels_out == i] = 1

            prod = np.multiply(mask, ptarray)

            num_nonzero_voxels = len(np.nonzero(mask)[0])

            lesion_suvmeans.append(np.sum(prod)/num_nonzero_voxels)

        tlg = np.sum(np.multiply(lesion_mtvs, lesion_suvmeans))

        return tlg

#%%

def calculate_patient_level_dissemination(

    maskarray: np.ndarray,

    spacing: tuple

) -> np.float64:

    """Function to return the tumor dissemination (Dmax) using 3D segmentation mask

    Dmax = max possible distance between any two foreground voxels in a patient;

    these two voxels can come form the same lesions (in case of one lesion) 

    or from different lesions (in case of multiple lesions) 

    Args:

        maskarray (np.ndarray): numpy array for 3D mask image

    Returns:

        np.float64: dissemination value in cm

    """

    maskarray = maskarray.astype(np.int8)

    nonzero_voxels = np.argwhere(maskarray == 1)

    distances = np.sqrt(np.sum(((nonzero_voxels[:, None] - nonzero_voxels) * spacing)**2, axis=2))

    farthest_indices = np.unravel_index(np.argmax(distances), distances.shape)

    dmax = distances[farthest_indices]/10  # converting to cm

    del maskarray 

    del nonzero_voxels

    del distances 

    return dmax 

#%%

def calculate_patient_level_dice_score(

    gtarray: np.ndarray,

    predarray: np.ndarray, 

) -> np.float64:

    """Function to return the Dice similarity coefficient (Dice score) between

    2 segmentation masks (containing 0s for background and 1s for lesions/tumors)

    Args:

        maskarray_1 (np.ndarray): numpy ndarray for the first mask

        maskarray_2 (np.ndarray): numpy ndarray for the second mask

    Returns:

        np.float64: Dice score

    """

    dice_score = 2.0*np.sum(predarray[gtarray == 1])/(np.sum(gtarray) + np.sum(predarray))

    return dice_score

#%%

def calculate_patient_level_iou(

    gtarray: np.ndarray,

    predarray: np.ndarray, 

) -> np.float64:

    """Function to return the Intersection-over-Union (IoU) between

    2 segmentation masks (containing 0s for background and 1s for lesions/tumors)

    Args:

        maskarray_1 (np.ndarray): numpy ndarray for the first mask

        maskarray_2 (np.ndarray): numpy ndarray for the second mask

    Returns:

        np.float64: Dice score

    """

    intersection = np.sum(predarray[gtarray == 1])

    union = np.sum(gtarray) + np.sum(predarray) - intersection

    iou = intersection/union

    return iou

def calculate_patient_level_intersection(

    gtarray: np.ndarray,

    predarray: np.ndarray, 

) -> np.float64:

    """Function to return the Intersection etween

    2 segmentation masks (containing 0s for background and 1s for lesions/tumors)

    Args:

        maskarray_1 (np.ndarray): numpy ndarray for the first mask

        maskarray_2 (np.ndarray): numpy ndarray for the second mask

    Returns:

        np.float64: Dice score

    """

    intersection = np.sum(predarray[gtarray == 1])

    return intersection

#%%

def calculate_patient_level_false_positive_volume(

    gtarray: np.ndarray,

    predarray: np.ndarray,

    spacing: tuple

) -> np.float64:

    # compute number of voxels of false positive connected components in prediction mask

    pred_connected_components = cc3d.connected_components(predarray, connectivity=18)

    false_positive = 0

    for idx in range(1,pred_connected_components.max()+1):

        comp_mask = np.isin(pred_connected_components, idx)

        if (comp_mask*gtarray).sum() == 0:

            false_positive += comp_mask.sum()

    voxel_volume_cc = np.prod(spacing)/1000

    return false_positive*voxel_volume_cc

#%%

def calculate_patient_level_false_negative_volume(

    gtarray: np.ndarray,

    predarray: np.ndarray,

    spacing: tuple

) -> np.float64:

    # compute number of voxels of false negative connected components (of the ground truth mask) in the prediction mask

    gt_connected_components = cc3d.connected_components(gtarray, connectivity=18)

    false_negative = 0

    for idx in range(1,gt_connected_components.max()+1):

        comp_mask = np.isin(gt_connected_components, idx)

        if (comp_mask*predarray).sum() == 0:

            false_negative += comp_mask.sum()

    voxel_volume_cc = np.prod(spacing)/1000

    return false_negative*voxel_volume_cc

# %%

def is_suvmax_detected(

    gtarray: np.ndarray,

    predarray: np.ndarray,

    ptarray: np.ndarray,

) -> bool:

    prod = np.multiply(gtarray, ptarray)

    max_index = np.unravel_index(np.argmax(prod), prod.shape)

    if predarray[max_index] == 1:

        return True

    else:

        return False

def calculate_patient_level_tp_fp_fn(

    gtarray: np.ndarray,

    predarray: np.ndarray,

    criterion: str,

    threshold: np.float64 = None,

    ptarray: np.ndarray = None,

) -> (int, int, int):

    """Calculate patient-level TP, FP, and FN (for detection based metrics)

    via 3 criteria:

    criterion1: A predicted lesion is TP if any one of it's foreground voxels 

    overlaps with GT foreground. A predicted lesions that doesn't overlap with any 

    GT foreground is FP. As soon as a lesion is predicted as TP, it is removed

    from the set of GT lesions. The lesions that remain in the end in the GT lesions

    are FN. `criterion1` is the weakest detection criterion.

    criterion2: A predicted lesion is TP if more than `threshold`% of it's volume 

    overlaps with foreground GT. A predicted lesion is FP if it overlap fraction

    with foreground GT is between 0% and `threshold`%. As soon as a lesion is 

    predicted as TP, it is removed from the set of GT lesions. The lesions that 

    remain in the end in the GT lesions are FN. `criterion2` can be hard or weak 

    criterion based on the value of `threshold`.

    criterion3: A predicted lesion is TP if it overlaps with one the the GT lesion's 

    SUVmax voxel, hence this criterion requires the use of PET data (`ptarray`). A 

    predicted lesion that doesn't overlap with any GT lesion's SUVmax voxel is 

    considered FP. As soon as a lesion is predicted as TP, it is removed from the 

    set of GT lesions. The lesions that remain in the end in the GT lesions are FN. 

    `criterion3` is likely an easy criterion since a network is more likely to segment 

    high(er)-uptake regions`.

    Args:

        int (_type_): _description_

        int (_type_): _description_

        gtarray (_type_, optional): _description_. Defaults to None, ptarray: np.ndarray = None, )->(int.

    """

    gtarray_labeled_mask, num_lesions_gt = cc3d.connected_components(gtarray, connectivity=18, return_N=True)

    predarray_labeled_mask, num_lesions_pred = cc3d.connected_components(predarray, connectivity=18, return_N=True)

    gt_lesions_list = list(np.arange(1, num_lesions_gt+1))

    #initial values for TP, FP, FN

    TP = 0

    FP = 0 

    FN = num_lesions_gt 

    if criterion == 'criterion1':

        FN = 0 # for this criterion we are counting the number of FPs from 0 onwards, hence the reassignment

        for i in range(1, num_lesions_pred+1):

            pred_lesion_mask = np.where(predarray_labeled_mask == i, 1, 0)

            if np.any(pred_lesion_mask & (gtarray_labeled_mask > 0)):

                TP += 1

            else:

                FP += 1

        for j in range(1, num_lesions_gt+1):

            gt_lesion_mask = np.where(gtarray_labeled_mask == j, 1, 0)

            if not np.any(gt_lesion_mask & (predarray_labeled_mask > 0)):

                FN += 1

    elif criterion == 'criterion2':

        for i in range(1, num_lesions_pred+1):

            max_iou = 0

            match_gt_lesion = None 

            pred_lesion_mask = np.where(predarray_labeled_mask == i, 1, 0)

            for j in range(1, num_lesions_gt+1):

                gt_lesion_mask = np.where(gtarray_labeled_mask == j, 1, 0)

                iou = calculate_patient_level_iou(gt_lesion_mask, pred_lesion_mask)

                if iou > max_iou:

                    max_iou = iou

                    match_gt_lesion = j

            if max_iou >= threshold:

                TP += 1

                gt_lesions_list.remove(match_gt_lesion)

            else:

                FP += 1

        FN = len(gt_lesions_list)

    elif criterion == 'criterion3':

        for i in range(1, num_lesions_pred+1):

            max_iou = 0

            match_gt_lesion = None

            pred_lesion_mask = np.where(predarray_labeled_mask == i, 1, 0)

            for j in range(1, num_lesions_gt+1):

                gt_lesion_mask = np.where(gtarray_labeled_mask == j, 1, 0)

                iou = calculate_patient_level_iou(gt_lesion_mask, pred_lesion_mask)

                if iou > max_iou:

                    max_iou = iou 

                    match_gt_lesion = j

            # match_gt_lesion has been defined with has the maximum iou with pred lesion i

            arr_gt_lesion = np.where(gtarray_labeled_mask == match_gt_lesion, 1, 0)

            if is_suvmax_detected(arr_gt_lesion, pred_lesion_mask, ptarray):

                TP += 1

                gt_lesions_list.remove(match_gt_lesion)

            else:

                FP += 1

        FN = len(gt_lesions_list)

    else:

        print('Invalid criterion. Choose between criterion1, criterion2, or criterion3')

        return 

    return TP, FP, FN