# example of using saved cycleGAN models for image translation
#based on https://machinelearningmastery.com/cyclegan-tutorial-with-keras/
from keras.models import load_model
import numpy as np
import tensorflow_addons as tfa
from scipy.ndimage import zoom
from tqdm import tqdm
import warnings
warnings.filterwarnings("ignore")
from huggingface_hub import hf_hub_download
from skimage.morphology import binary_erosion, binary_dilation
from skimage import draw


def predict_mask(image, dim_x, dim_y, dim_z, _resize=True, norm_=True, mode_='test', patch_size=(64,128,128,1), _step=64, _step_z=32, _patch_size_z=64):

    cust={'InstanceNormalization': tfa.layers.InstanceNormalization}
    #load the model
    # Download the model from Hugging Face Model Hub
    model_dir = hf_hub_download(repo_id="Hemaxi/3DCycleGAN", filename="CycleGANVesselSegmentation.h5")
    model_BtoA = load_model(model_dir, cust)

    print('Mode: {}'.format(mode_))

    _patch_size = patch_size[1]
    _nbslices = patch_size[0]

    perceqmin = 1
    perceqmax = 99

    image = ((image/(np.max(image)))*255).astype('uint8')

    print('Image Shape: {}'.format(image.shape))
    print('----------------------------------------')

    initial_image_x = np.shape(image)[0]
    initial_image_y = np.shape(image)[1]
    initial_image_z = np.shape(image)[2]

    #percentile equalization
    if norm_:
        minval = np.percentile(image, perceqmin) 
        maxval = np.percentile(image, perceqmax)
        image = np.clip(image, minval, maxval)
        image = (((image - minval) / (maxval - minval)) * 255).astype('uint8')

    if _resize:
        image = zoom(image, (dim_x/0.333, dim_y/0.333, dim_z/0.5), order=3, mode='nearest')
        image = ((image/np.max(image))*255.0).astype('uint8')


    #image size
    size_y = np.shape(image)[0]
    size_x = np.shape(image)[1]
    size_depth = np.shape(image)[2]
    aux_sizes_or = [size_y, size_x, size_depth]
    

    #patch size
    new_size_y = int((size_y/_patch_size) + 1) * _patch_size
    new_size_x = int((size_x/_patch_size) + 1) * _patch_size
    new_size_z = int((size_depth/_patch_size_z) + 1) * _patch_size_z
    aux_sizes = [new_size_y, new_size_x, new_size_z]
    
    ## zero padding
    aux_img = np.random.randint(1,50,(aux_sizes[0], aux_sizes[1], aux_sizes[2]))
    aux_img[0:aux_sizes_or[0], 0:aux_sizes_or[1],0:aux_sizes_or[2]] = image
    image = aux_img
    del aux_img
        
    final_mask_foreground = np.zeros((np.shape(image)[0], np.shape(image)[1], np.shape(image)[2]))
    final_mask_background = np.zeros((np.shape(image)[0], np.shape(image)[1], np.shape(image)[2]))
    final_mask_background = final_mask_background.astype('uint8')
    final_mask_foreground = final_mask_foreground.astype('uint8')
    

    total_iterations = int(image.shape[0]/_patch_size)

    with tqdm(total=total_iterations) as pbar:
        i=0
        while i+_patch_size<=image.shape[0]:
            j=0
            while j+_patch_size<=image.shape[1]:
                k=0
                while k+_patch_size_z<=image.shape[2]:
                
                    B_real = np.zeros((1,_nbslices,_patch_size,_patch_size,1),dtype='float32')
                    _slice = image[i:i+_patch_size, j:j+_patch_size, k:k+_patch_size_z]
                    
                    _slice = _slice.transpose(2,0,1)
                    _slice = np.expand_dims(_slice, axis=-1)
        
                    B_real[0,:]=(_slice-127.5) /127.5   
        
                    A_generated  = model_BtoA.predict(B_real)
        
                    A_generated = (A_generated + 1)/2 #from [-1,1] to [0,1]
        
                    A_generated = A_generated[0,:,:,:,0]
                    A_generated = A_generated.transpose(1,2,0)
        
                    #print(np.unique(A_generated))
                    A_generated = (A_generated>0.5)*1
        
                    A_generated = A_generated.astype('uint8')
        
                    final_mask_foreground[i:i+_patch_size, j:j+_patch_size, k:k+_patch_size_z] = final_mask_foreground[i:i+_patch_size, j:j+_patch_size, k:k+_patch_size_z] + A_generated
                    final_mask_background[i:i+_patch_size, j:j+_patch_size, k:k+_patch_size_z] = final_mask_background[i:i+_patch_size, j:j+_patch_size, k:k+_patch_size_z] + (1-A_generated)
                    
                    k=k+_step_z
                j=j+_step
            i=i+_step
            pbar.update(1)


    del _slice
    del A_generated
    del B_real

    final_mask = (final_mask_foreground>=final_mask_background)*1

    image = image[0:aux_sizes_or[0], 0:aux_sizes_or[1],0:size_depth]
    print('Image Shape: {}'.format(image.shape))
    print('----------------------------------------')

    final_mask = final_mask[0:aux_sizes_or[0], 0:aux_sizes_or[1],0:aux_sizes_or[2]]


    if _resize:
        final_mask = zoom(final_mask, (0.333/dim_x, 0.333/dim_y, 0.5/dim_z), order=3, mode='nearest')
        final_mask = (final_mask*255.0).astype('uint8')

        final_size_x = np.shape(final_mask)[0]
        final_size_y = np.shape(final_mask)[1]
        final_size_z = np.shape(final_mask)[2]

        aux_mask = np.zeros((initial_image_x, initial_image_y, initial_image_z)).astype('uint8')
        aux_mask[0:min(initial_image_x, final_size_x),0:min(initial_image_y, final_size_y),0:min(initial_image_z, final_size_z)] = final_mask[0:min(initial_image_x, final_size_x),0:min(initial_image_y, final_size_y),0:min(initial_image_z, final_size_z)]

        final_mask = aux_mask.copy()


    print('Mask Shape: {}'.format(final_mask.shape))
    print('----------------------------------------')
    final_mask = final_mask/np.max(final_mask)
    final_mask = final_mask*255.0
    final_mask = final_mask.astype('uint8')


    #closing operation to fill small holes
    mask = final_mask
    mask[mask!=0] = 1
    mask = mask.astype('uint8')
    
    ellipsoid = draw.ellipsoid(9,9,3, spacing=(1,1,1), levelset=False)
    ellipsoid = ellipsoid.astype('uint8')
    ellipsoid = ellipsoid[1:-1,1:-1,1:-1]
    
    #perform closing operation on the mask
    dil = binary_dilation(mask, ellipsoid)
    closed_mask = binary_erosion(dil, ellipsoid)
    closed_mask = (closed_mask*255.0).astype('uint8')
    
    return closed_mask