import os
import sys
import tempfile

import numpy as np
import vreg

from miblab.data import zenodo_fetch
from miblab.data import clear_cache_datafiles


if sys.version_info < (3, 9):
    # importlib.resources either doesn't exist or lacks the files()
    # function, so use the PyPI version:
    import importlib_resources
else:
    # importlib.resources has files(), so use that:
    import importlib.resources as importlib_resources

try:
    from monai.networks.nets.unetr import UNETR
    from monai.inferers import sliding_window_inference
    monai_installed = True
except ImportError:
    monai_installed = False

try:
    import torch
    torch_installed = True
except ImportError:
    torch_installed = False

try:
    import nibabel as nib
    nib_installed = True
except ImportError:
    nib_installed = False

try:
    from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor
    nnunetv2 = True
except ImportError:
    nnunetv2 = False

try:
    import scipy.ndimage as ndi
    scipy_installed = True
except ImportError:
    scipy_installed = False


def kidney_pc_dixon_unetr(input_vol, device=None, overlap=0.3, postproc=True, clear_cache = False, verbose=False):

    """
    Segment individual kidneys on post-contrast Dixon images.

    This requires 4-channel input data with out-phase images, 
    in-phase images, water maps, and fat maps.

    Args:
        input_vol (vreg.Volume3D): A 4D vreg volume of shape 
            [x, y, z, contrast] representing the input medical image 
            volume. The last index must contain out-phase, in-phase, 
            water and fat images, in that order.
        device (str):
            processor on which to deply the computation. If this is not 
            provided, this defaults to "cuda" (if this is available) and 
            "cpu" otherwise.
        overlap (float): only valid for model = 'unetr' defines the amount of overlap between 
            adjacent sliding window patches during inference. A 
            higher value (e.g., 0.5) improves prediction smoothness 
            at patch borders but increases computation time.
        postproc (bool): If True, applies post-processing to select 
            the largest connected component from the UNETR output 
            for each kidney mask
        clear_cache: If True, the downloaded pth file is removed 
            again after running the inference.
        verbose (bool): If True, prints logging messages.

    Returns:
        vreg.Volume3D: 
            A volume with a kidneys as label array.
    """

    if not torch_installed:
        raise ImportError(
            'torch is not installed. Please install it with "pip install torch".'
            'To install all dlseg options at once, install miblab as pip install miblab[dlseg].'
        )
    if not monai_installed:
        raise ImportError(
            'monai is not installed. Please install it with "pip install monai".'
            'To install all dlseg options at once, install miblab as pip install miblab[dlseg].'
        )
    if not scipy_installed:
        raise ImportError(
            'scipy is not installed. Please install it with "pip install scipy".'
            'To install all dlseg options at once, install miblab as pip install miblab[dlseg].'
        )


    MODEL = 'UNETR_kidneys_v2.pth'
    MODEL_DOI = "15521814"

    input_array = input_vol.values

    # Format input array
    # ------------------
    # TODO: make this consistent for future training
    # swap the last two indices because UNETR is trained 
    # with opposed-phase/in-phase/fat/water
    input_array = input_array[..., [0, 1, 3, 2]]

    # flip
    if not input_vol.is_right_handed:
        input_array = np.flip(input_array, axis=2)

    # Rescale to required matrix size
    old_shape = input_array.shape
    new_shape = (320, 320, 144, 4)
    zoom_factors = [n / o for n, o in zip(new_shape, old_shape)]
    input_array = ndi.zoom(input_array, zoom_factors, order=1) 

    # from (x,y,z,c) to (c,y,x,z)
    input_array = np.transpose(input_array, (3, 0, 1, 2)) 

    # Normalize data
    input_array_out   = (input_array[0,...]-np.average(input_array[0,...]))/np.std(input_array[0,...])
    input_array_in    = (input_array[1,...]-np.average(input_array[1,...]))/np.std(input_array[1,...])
    input_array_water = (input_array[2,...]-np.average(input_array[2,...]))/np.std(input_array[2,...])
    input_array_fat   = (input_array[3,...]-np.average(input_array[3,...]))/np.std(input_array[3,...])

    input_array = np.stack((input_array_out, input_array_in, input_array_water, input_array_fat), axis=0)
    # Convert to NCHW[D] format: (1,c,y,x,z)
    # NCHW[D] stands for: batch N, channels C, height H, width W, depth D
    input_array = input_array.transpose(0,2,1,3) # from (x,y,z) to (y,x,z)
    input_array = np.expand_dims(input_array, axis=(0))
    input_array = input_array.astype(np.float32)

    if verbose:
        print('(Down)loading model..')

    weights_dir = importlib_resources.files('miblab.datafiles')
    weights_path = zenodo_fetch(MODEL, weights_dir, MODEL_DOI)

    if verbose:
        print('Applying model to data..')

    # Setup device
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    if device is None:
        device_str = "cuda" if torch.cuda.is_available() else "cpu"
    else:
        device_str = device
    device = torch.device(device_str)

    # Define model architecture
    model = UNETR(
        in_channels=4,
        out_channels=3, # BACKGROUND, RIGHT KIDNEY (left on image), LEFT KIDNEY (right on image)
        img_size=(80, 80, 80),
        feature_size=16,
        hidden_size=768,
        mlp_dim=3072,
        num_heads=12,
        proj_type="perceptron",
        norm_name="instance",
        res_block=True,
        dropout_rate=0.0,
    )
    model.to(device)

    # Convert to tensor
    input_tensor = torch.tensor(input_array).to(device)

    # Load model weights
    weights = torch.load(weights_path, map_location=device, weights_only=True)
    model.load_state_dict(weights)
    model.eval() 

    with torch.no_grad():
        output_tensor = sliding_window_inference(input_tensor, (80,80,80), 4, model, overlap=overlap, device=device_str, progress=True) 

    if verbose:
        print('Post-processing results...')

    # From probabilities for each channel to label image
    output_tensor = torch.argmax(output_tensor, dim=1)

    # Convert to numpy
    output_array = output_tensor.numpy(force=True)[0,:,:,:]
        
    # Transpose to original shape
    output_array = output_array.transpose(1,0,2) #from (y,x,z) to (x,y,z)

    # Rescale to original matrix size
    zoom_factors = [n / o for n, o in zip(old_shape[:3], new_shape[:3])]
    output_array = ndi.zoom(output_array, zoom_factors, order=1) 

    # flip back
    if not input_vol.is_right_handed:
        output_array = np.flip(output_array, axis=2)
    
    if postproc == True:
        left_kidney, right_kidney = _kidney_masks(output_array)
    else:
        left_kidney=output_array[output_array == 2]
        left_kidney[left_kidney==2]=1
        right_kidney=output_array[output_array == 1]

    kidneys = (left_kidney + 2*right_kidney).astype(np.int16)

    if clear_cache:
        if verbose:
            print('Deleting downloaded files...')
        clear_cache_datafiles(weights_dir)

    return vreg.volume(kidneys, input_vol.affine) 


def kidney_pc_dixon(input_array, device=None, postproc=True, clear_cache = False, verbose=False):

    """
    Segment individual kidneys on post-contrast Dixon images.

    This requires 4-channel input data with out-phase images, 
    in-phase images, water maps, and fat maps.

    Args:
        input_array (numpy.ndarray): A 4D numpy array of shape 
            [x, y, z, contrast] representing the input medical image 
            volume. The last index must contain out-phase, in-phase, 
            water and fat images, in that order.
        device (str):
            processor on which to deply the computation. If this is not 
            provided, this defaults to "cuda" (if this is available) and 
            "cpu" otherwise.
        postproc (bool): If True, applies post-processing to select 
            the largest connected component from the UNETR output 
            for each kidney mask
        clear_cache: If True, the downloaded pth file is removed 
            again after running the inference.
        verbose (bool): If True, prints logging messages.

    Returns:
        numpy.ndarray: 
            Label image

    """

    if not torch_installed:
        raise ImportError(
            'torch is not installed. Please install it with "pip install torch".'
            'To install all dlseg options at once, install miblab as pip install miblab[dlseg].'
        )
    if not monai_installed:
        raise ImportError(
            'totalsegmentator is not installed. Please install it with "pip install totalsegmentator".'
            'To install all dlseg options at once, install miblab as pip install miblab[dlseg].'
        )
    if not nib_installed:
        raise ImportError(
            'nibabel is not installed. Please install it with "pip install nibabel".'
            'To install all dlseg options at once, install miblab as pip install miblab[dlseg].'
        )
    if not scipy_installed:
        raise ImportError(
            'scipy is not installed. Please install it with "pip install scipy".'
            'To install all dlseg options at once, install miblab as pip install miblab[dlseg].'
        )
  
    MODEL = 'nnunet_kidneys_v2.zip'
    MODEL_DOI = "15328218"

    if verbose:
        print('(Down)loading model..')

    weights_dir = importlib_resources.files('miblab.datafiles')
    weights_path = zenodo_fetch(MODEL, weights_dir, MODEL_DOI, extract=True)

    if verbose:
        print('Applying model to data..')

    # Setup device
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    if device is None:
        device_str = "cuda" if torch.cuda.is_available() else "cpu"
    else:
        device_str = device
    device = torch.device(device_str)

    # Setup device
    predictor = nnUNetPredictor(
        tile_step_size=0.75,
        use_gaussian=True,
        use_mirroring=False,
        perform_everything_on_device=True,
        device=device,
        verbose=False,
        verbose_preprocessing=False,
        allow_tqdm=True
    )

    # direct to "nnUNetTrainer__nnUNetPlans__3d_fullres"
    nested_folder = os.path.join(weights_path, "nnUNetTrainer__nnUNetPlans__3d_fullres")

    # Setup the predictor
    predictor.initialize_from_trained_model_folder(
        nested_folder,
        use_folds='all',
        checkpoint_name='checkpoint_best.pth'
    )

    with tempfile.TemporaryDirectory() as temp_dir:

        # Generate temp folder for intermediate data
        temp_folder_results = os.path.join(temp_dir,"temp_results")
        temp_folder_data_to_test = os.path.join(temp_dir,"temp_results",'data_to_test')
        os.makedirs(temp_folder_data_to_test, exist_ok=True)

        # Save arrays as nifti (.nii)
        affine = np.eye(4)
        nii_out_ph = nib.Nifti1Image(input_array[...,0], affine)
        nib.save(nii_out_ph, os.path.join(temp_folder_data_to_test, 'Dixon_999_0000.nii.gz'))
        nii_in_ph = nib.Nifti1Image(input_array[...,1], affine)
        nib.save(nii_in_ph, os.path.join(temp_folder_data_to_test, 'Dixon_999_0001.nii.gz'))
        nii_water = nib.Nifti1Image(input_array[...,2], affine)
        nib.save(nii_water, os.path.join(temp_folder_data_to_test, 'Dixon_999_0002.nii.gz'))
        nii_fat = nib.Nifti1Image(input_array[...,3], affine)
        nib.save(nii_fat, os.path.join(temp_folder_data_to_test, 'Dixon_999_0003.nii.gz'))

        # Infer kidney masks
        predictor.predict_from_files(
            temp_folder_data_to_test,
            temp_folder_results, 
            save_probabilities=False, 
            overwrite=True
        )

        # Load the NIfTI file
        nifti_file = nib.load(os.path.join(temp_folder_results,'Dixon_999.nii.gz'))
        output_array = nifti_file.get_fdata()
    
    if postproc == True:
        left_kidney, right_kidney = _kidney_masks(output_array)
    else:
        left_kidney=output_array[output_array == 2]
        left_kidney[left_kidney==2]=1
        right_kidney=output_array[output_array == 1]

    kidneys = (left_kidney + 2*right_kidney).astype(np.int16)

    if clear_cache:
        if verbose:
            print('Deleting downloaded files...')
        clear_cache_datafiles(temp_dir)

    return kidneys  


def _largest_cluster(array:np.ndarray)->np.ndarray:
    """Given a mask array, return a new mask array containing only the largesr cluster.

    Args:
        array (np.ndarray): mask array with values 1 (inside) or 0 (outside)

    Returns:
        np.ndarray: mask array with only a single connect cluster of pixels.
    """
    # Label all features in the array
    label_img, cnt = ndi.label(array)
    # Find the label of the largest feature
    labels = range(1,cnt+1)
    size = [np.count_nonzero(label_img==l) for l in labels]
    max_label = labels[size.index(np.amax(size))]
    # Return a mask corresponding to the largest feature
    return label_img==max_label

def _kidney_masks(output_array:np.ndarray)->tuple:
    """Extract kidney masks from the output array of the UNETR

    Args:
        output_array (np.ndarray): 3D numpy array (x,y,z) with integer labels (0=background, 1=right kidney, 2=left kidney)

    Returns:
        tuple: A tuple of 3D numpy arrays (left_kidney, right_kidney) with masks for the kidneys.
    """
    left_kidney = _largest_cluster(output_array == 2)
    right_kidney = _largest_cluster(output_array == 1)

    return left_kidney, right_kidney