import numpy as np
import sys
import os
import glob
import nibabel as nib
import torch
#docker
#fpx = '/wmhseg_code/'
#inputDir = '/input'
#outputDir = '/output'

#local
main_folder_path = '../Data/MS2017b/'
fpx = './'
inputDir = 'input/'
outputDir = 'output/'

#PARAMS
useGPU = 0
gpu0 = 0
patch_size = 60
extra_patch = 5
model_paths = [fpx + 'analysis/models/EXP3D_1x1x1x1_0_0_dice_1_best.pth']
weights = [1]
#EXP3D_1x1x1x1_0_0_dice_1_best.pth
#EXP3D_2x2x2x2_0_0_dice_1_best.pth
#EXP3D_1x1x1x1_0_1_dice_1_best.pth
#EXP3D_1x1x1x1_1_1_dice_1_best.pth
##EXP3D_1x1x1x1_0_1_dice_1_best.pth (with ep = 16)
sys.path.append(fpx + 'utils/')
sys.path.append(fpx + 'architectures/deeplab_3D/')
sys.path.append(fpx + 'architectures/unet_3D/')
sys.path.append(fpx + 'architectures/hrnet_3D/')
sys.path.append(fpx + 'architectures/experiment_nets_3D/')
sys.path.append('utils/')

import deeplab_resnet_3D
import unet_3D
import highresnet_3D
import exp_net_3D

import augmentations as AUG
import normalizations as NORM
import resizeScans as RS
import evalF as EF
import evalFP as EFP
import PP
import torch

#step 1: read image from input folder
#step 2: resize image to 200x200x100 + apply normalizations
#step 3: make prediction by patches (with augmentations)
#step 4: save prediction to output folder
#step 5: resize prediction back to original size of image


img_path = os.path.join(inputDir, 'FLAIR.nii.gz')
img_path_rs = os.path.join(outputDir, 'FLAIR_rs.nii.gz')

wmh_path_rs = os.path.join(outputDir, 'wmh_rs.nii.gz')
wmh_path = os.path.join(outputDir, 'result.nii.gz')

old_size = PP.numpyFromScan(img_path).shape

new_size = [200,200,100]
num_labels = 2

#convert scan to 200x200x100
RS.convertSize2(img_path, img_path_rs, new_size)
#get the affine value
affine_rs = nib.load(img_path_rs).get_affine()

#normalize using histogram and variance normalization
RS.normalizeScan(img_path_rs, img_path_rs, main_folder_path=main_folder_path)

#read preprocessed img
img, affine = PP.numpyFromScan(img_path_rs, get_affine = True)
img = img.transpose((3,0,1,2))
img = img[np.newaxis, :]

print('Image ready')
print('Loading model')

out = None
for i, model_path in enumerate(model_paths): 
	f_name = model_path.split('/')[-1]
	isPriv = False

	#load model
	if 'EXP3D' in f_name:
		experiment = f_name.replace('EXP3D_', '').replace('.pth', '').split('_')
		experiment = '_'.join(experiment[0:3])
		dilation_arr, isPriv, withASPP = PP.getExperimentInfo(experiment)
		model = exp_net_3D.getExpNet(num_labels, dilation_arr, isPriv, NoLabels2 = 209, withASPP = withASPP)
	elif 'HR3D' in f_name:
		model = highresnet_3D.getHRNet(num_labels)
	elif 'DL3D' in f_name:
		model = deeplab_resnet_3D.Res_Deeplab(num_labels)
	elif 'UNET3D' in  f_name:
		model = unet_3D.UNet3D(1, num_labels)
	else:
		print('No model available for this .pth')
		sys.exit()

	if useGPU:
	    saved_state_dict = torch.load(model_path)
	else:
	    saved_state_dict = torch.load(model_path, map_location=lambda storage, loc: storage)
	model.load_state_dict(saved_state_dict)
	model.float()
	model.eval()
	print('Model ready')
	print('Predicting...')
	if not isinstance(out, np.ndarray):
		if isPriv:
			out = EFP.testPredict(img, model, num_labels, 209, 1, gpu0, useGPU, stride = 50, patch_size = 60, test_augm = False, extra_patch = extra_patch, get_soft = True)
		else:
			out = EF.testPredict(img, model, num_labels, 1, gpu0, useGPU, stride = 50, patch_size = 60, test_augm = False, extra_patch = extra_patch, get_soft = True)
	else:
		if isPriv:
			out += EFP.testPredict(img, model, num_labels, 209, 1, gpu0, useGPU, stride = 50, patch_size = 60, test_augm = False, extra_patch = extra_patch, get_soft = True)
		else:
			out += EF.testPredict(img, model, num_labels, 1, gpu0, useGPU, stride = 50, patch_size = 60, test_augm = False, extra_patch = extra_patch, get_soft = True)

out /= float(len(model_paths))
out = np.argmax(out, axis = 0)
#remove batch and label dimension
out = out.squeeze()
print('Prediction complete')
print('Saving...')
#save output
PP.saveScan(out.astype(np.float64), affine_rs, wmh_path_rs)

#resize output to original input size and save (this is our final result)
d = RS.convertSize2(wmh_path_rs, wmh_path, old_size, interpolation = 'nearest')

#read the image and save it with same affine and header as original FLAIR image
print('Saving final wmh file')
orig_flair = nib.load(img_path)
wmh_final = nib.load(wmh_path).get_data()
PP.saveScan(wmh_final, orig_flair.get_affine(), wmh_path, header =orig_flair.header)
print('Done')