import os

from Segmenation import BreastSeg,BreastTumor

from ImageProcessing import read_image,Norm_Zscore,imgnorm,save_image,labeling_seg

import argparse

import torch

try:

    torch._utils._rebuild_tensor_v2

except AttributeError:

    def _rebuild_tensor_v2(storage, storage_offset, size, stride, requires_grad, backward_hooks):

        tensor = torch._utils._rebuild_tensor(storage, storage_offset, size, stride)

        tensor.requires_grad = requires_grad

        tensor._backward_hooks = backward_hooks

        return tensor

    torch._utils._rebuild_tensor_v2 = _rebuild_tensor_v2

parser = argparse.ArgumentParser(description='Tumor Segmentation from DCE-MRI')

parser.add_argument('--cuda', type=int, default='1', required=False, help='Run in GPU')

parser.add_argument('--pre', type=str, default='Data/Img1_pre.nii.gz', required=False, help='Image path for pre-constrast image')

parser.add_argument('--post', type=str, default='Data/Img1_post.nii.gz', required=False, help='Image path for post-constrast image')

parser.add_argument('--outfolder',type=str,default='Results',required=False,help='Folder for saving results')

opt = parser.parse_args()

print(opt)

if opt.cuda and not torch.cuda.is_available():

    raise Exception("No GPU found, please run without --cuda 0")

savepath = opt.outfolder

if not os.path.exists(savepath):

    os.makedirs(savepath)

# read images and output spacing information   

print('Reading images')      

img_pre, img_post, img_sub, scale_subject = read_image(opt)

# perform the intensity normalization 1) remover outliers 2) Z-score normalization 

# spatial normalization will be done before segmentation in the functions of BrestSeg and BreastTumor

print('Performing image normalization')

img_pre  = Norm_Zscore(imgnorm(img_pre))

img_post = Norm_Zscore(imgnorm(img_post))

img_sub  = Norm_Zscore(imgnorm(img_sub))

# read models

# the weights of the models will be loaded in the function of BrestSeg and BreastTumor

from Models_3D import ModelBreast,ModelTumor

if opt.cuda:

    model_breast = ModelBreast(1,1).cuda()

else:

    model_breast = ModelBreast(1,1)

if opt.cuda:

    model_tumor = ModelTumor(3,1).cuda()

else:

    model_tumor = ModelTumor(3,1)

# perform image segmentation    

print('Performing image segmentation')

breast_mask = BreastSeg(img_pre,scale_subject,model_breast,opt)

prob_1st, seg_2nd = BreastTumor(img_sub,img_post,breast_mask,scale_subject,model_tumor,model_tumor,opt)

# perform labeling

labeled_image = labeling_seg(seg_2nd)

# save images

print('Saving segmentations')

save_image(breast_mask,opt,'breast_mask.nii.gz')

save_image(prob_1st,opt,'prob_1st.nii.gz')

save_image(labeled_image,opt,'seg_2nd.nii.gz')

print('Please check the segmentation results in the result folder ---> "Results"')