import SimpleITK as sitk
import tensorflow as tf
from tensorflow.keras.models import load_model
from focal_loss import BinaryFocalLoss
import numpy as np
import multiprocessing
from functools import partial
import os 
from os import path
from tqdm import tqdm
import argparse

from sfransen.utils_quintin import * 
from sfransen.DWI_exp.helpers import *
from sfransen.DWI_exp.preprocessing_function import preprocess  
from sfransen.DWI_exp.callbacks import dice_coef
#from sfransen.FROC.blob_preprocess import *
from sfransen.FROC.cal_froc_from_np import *
from sfransen.load_images import load_images_parrallel
from sfransen.DWI_exp.losses import weighted_binary_cross_entropy


parser = argparse.ArgumentParser(
    description='Calculate the froc metrics and store in froc_metrics.yml')
parser.add_argument('-experiment',
    help='Title of experiment')                
parser.add_argument('--series', '-s', 
    metavar='[series_name]', required=True, nargs='+',
    help='List of series to include')
parser.add_argument('-fold',
    default='',
    help='List of series to include')
args = parser.parse_args()

# if __name__ = '__main__':
# bovenstaande nodig om fork probleem op te lossen (windows cs linux)

######## CUDA ################
os.environ["CUDA_VISIBLE_DEVICES"] = "2"

######## constants #############
fold = args.fold
SERIES = args.series
series_ = '_'.join(args.series)
EXPERIMENT = args.experiment

MODEL_PATH = f'./../train_output/{EXPERIMENT}_{series_}_{fold}/models/{EXPERIMENT}_{series_}_{fold}.h5'
YAML_DIR = f'./../train_output/{EXPERIMENT}_{series_}_{fold}'
IMAGE_DIR = f'./../train_output/{EXPERIMENT}_{series_}_{fold}'

# MODEL_PATH = f'./../train_output/{EXPERIMENT}_{series_}/models/{EXPERIMENT}_{series_}.h5'
# YAML_DIR = f'./../train_output/{EXPERIMENT}_{series_}'
# IMAGE_DIR = f'./../train_output/{EXPERIMENT}_{series_}'

DATA_DIR = "./../data/Nijmegen paths/"
TARGET_SPACING = (0.5, 0.5, 3) 
INPUT_SHAPE = (192, 192, 24, len(SERIES))
IMAGE_SHAPE = INPUT_SHAPE[:3]

DATA_SPLIT_INDEX = read_yaml_to_dict(f'./../data/Nijmegen paths/train_val_test_idxs_{fold}.yml')
# DATA_SPLIT_INDEX = read_yaml_to_dict(f'./../data/Nijmegen paths/train_val_test_idxs.yml')
TEST_INDEX = DATA_SPLIT_INDEX['test_set0']

N_CPUS = 12


########## load images in parrallel ##############
print_(f"> Loading images into RAM...")

# # read paths from txt
# image_paths = {}
# for s in SERIES:
#     with open(path.join(DATA_DIR, f"{s}.txt"), 'r') as f:
#         image_paths[s] = [l.strip() for l in f.readlines()]
# with open(path.join(DATA_DIR, f"seg.txt"), 'r') as f:
#     seg_paths = [l.strip() for l in f.readlines()]
# num_images = len(seg_paths)

# # create pool of workers
# pool = multiprocessing.Pool(processes=N_CPUS)
# partial_images = partial(load_images_parrallel,
#                     seq = 'images',
#                     target_shape=IMAGE_SHAPE,
#                     target_space = TARGET_SPACING)
# partial_seg = partial(load_images_parrallel,
#                     seq = 'seg',
#                     target_shape=IMAGE_SHAPE,
#                     target_space = TARGET_SPACING)

# #load images
# images = []
# for s in SERIES:
#     image_paths_seq = image_paths[s]
#     image_paths_index = np.asarray(image_paths_seq)[TEST_INDEX][:5]
#     data_list = pool.map(partial_images,image_paths_index)
#     data = np.stack(data_list, axis=0)
#     images.append(data)

# images_list = np.transpose(images, (1, 2, 3, 4, 0))
# print('>>>>> size image_list nmr 1:', np.shape(images_list))

# #load segmentations
# seg_paths_index = np.asarray(seg_paths)[TEST_INDEX][:5]
# data_list = pool.map(partial_seg,seg_paths_index)
# segmentations = np.stack(data_list, axis=0)

########## test with old method #############
image_paths = {}

for s in SERIES:
    with open(path.join(DATA_DIR, f"{s}.txt"), 'r') as f:
        image_paths[s] = [l.strip() for l in f.readlines()]
with open(path.join(DATA_DIR, f"seg.txt"), 'r') as f:
    seg_paths = [l.strip() for l in f.readlines()]
num_images = len(seg_paths)

images = []
images_list = []
segmentations = []

# Read and preprocess each of the paths for each series, and the segmentations.
for img_idx in tqdm(range(len(TEST_INDEX))): #[:40]): #for less images
    # print('images number',[TEST_INDEX[img_idx]])
    img_s = {f'{s}': sitk.ReadImage(image_paths[s][TEST_INDEX[img_idx]], sitk.sitkFloat32) for s in SERIES}
    seg_s = sitk.ReadImage(seg_paths[TEST_INDEX[img_idx]], sitk.sitkFloat32)
    img_n, seg_n = preprocess(img_s, seg_s, 
        shape=IMAGE_SHAPE, spacing=TARGET_SPACING)
    for seq in img_n:
        images.append(img_n[f'{seq}'])
    images_list.append(images)
    images = []
    segmentations.append(seg_n)

images_list = np.transpose(images_list, (0, 2, 3, 4, 1))

print('>>>>> size image_list nmr 2:', np.shape(images_list), '. equal to: (5, 192, 192, 24, 3)?')

########### load module ##################
print(' >>>>>>> LOAD MODEL <<<<<<<<<')

dependencies = {
    'dice_coef': dice_coef,
    'weighted_cross_entropy_fn':weighted_binary_cross_entropy
}
reconstructed_model = load_model(MODEL_PATH, custom_objects=dependencies)
# reconstructed_model.summary(line_length=120)

# make predictions on all TEST_INDEX
print(' >>>>>>> START prediction <<<<<<<<<')
predictions_blur = reconstructed_model.predict(images_list, batch_size=1)

############# preprocess ################# 
# preprocess predictions by removing the blur and making individual blobs 
print('>>>>>>>> START preprocess')

def move_dims(arr):
    # UMCG numpy dimensions convention: dims = (batch, width, heigth, depth)
    # Joeran numpy dimensions convention: dims = (batch, depth, heigth, width)
    arr = np.moveaxis(arr, 3, 1) 
    arr = np.moveaxis(arr, 3, 2)
    return arr

# Joeran has his numpy arrays ordered differently.
predictions_blur = move_dims(np.squeeze(predictions_blur))
segmentations = move_dims(np.squeeze(segmentations))
predictions = [preprocess_softmax(pred, threshold="dynamic")[0] for pred in predictions_blur]

# Remove outer edges 
zeros = np.zeros(np.shape(predictions))
test = np.squeeze(predictions)[:,2:-2,2:190,2:190]
zeros[:,2:-2,2:190,2:190] = test
predictions = zeros

# perform Froc
metrics = evaluate(y_true=segmentations, y_pred=predictions)
dump_dict_to_yaml(metrics, YAML_DIR, "froc_metrics_focal_10_test", verbose=True)


############## save image as example #################
# save image nmr 2
img_s = sitk.GetImageFromArray(predictions_blur[2].squeeze())
sitk.WriteImage(img_s, f"{IMAGE_DIR}/predictions_blur_002_old.nii.gz")

img_s = sitk.GetImageFromArray(predictions[2].squeeze())
sitk.WriteImage(img_s, f"{IMAGE_DIR}/predictions_002_old.nii.gz")

img_s = sitk.GetImageFromArray(segmentations[2].squeeze())
sitk.WriteImage(img_s, f"{IMAGE_DIR}/segmentations_002_old.nii.gz")

img_s = sitk.GetImageFromArray(np.transpose(images_list[2,:,:,:,0].squeeze()))
sitk.WriteImage(img_s, f"{IMAGE_DIR}/t2_002_old.nii.gz")

# save image nmr 3
img_s = sitk.GetImageFromArray(predictions_blur[3].squeeze())
sitk.WriteImage(img_s, f"{IMAGE_DIR}/predictions_blur_003_old.nii.gz")

img_s = sitk.GetImageFromArray(predictions[3].squeeze())
sitk.WriteImage(img_s, f"{IMAGE_DIR}/predictions_003_old.nii.gz")

img_s = sitk.GetImageFromArray(segmentations[3].squeeze())
sitk.WriteImage(img_s, f"{IMAGE_DIR}/segmentations_003_old.nii.gz")

img_s = sitk.GetImageFromArray(np.transpose(images_list[3,:,:,:,0].squeeze()))
sitk.WriteImage(img_s, f"{IMAGE_DIR}/t2_003_old.nii.gz")

img_s = sitk.GetImageFromArray(np.transpose(images_list[3,:,:,:,1].squeeze()))
sitk.WriteImage(img_s, f"{IMAGE_DIR}/highb_003_old.nii.gz")

# save image nmr 3
img_s = sitk.GetImageFromArray(predictions_blur[4].squeeze())
sitk.WriteImage(img_s, f"{IMAGE_DIR}/predictions_blur_004_old.nii.gz")

img_s = sitk.GetImageFromArray(predictions[4].squeeze())
sitk.WriteImage(img_s, f"{IMAGE_DIR}/predictions_004_old.nii.gz")

img_s = sitk.GetImageFromArray(segmentations[4].squeeze())
sitk.WriteImage(img_s, f"{IMAGE_DIR}/segmentations_004_old.nii.gz")

img_s = sitk.GetImageFromArray(np.transpose(images_list[4,:,:,:,0].squeeze()))
sitk.WriteImage(img_s, f"{IMAGE_DIR}/t2_004_old.nii.gz")

img_s = sitk.GetImageFromArray(np.transpose(images_list[4,:,:,:,1].squeeze()))
sitk.WriteImage(img_s, f"{IMAGE_DIR}/highb_004_old.nii.gz")