from inspect import _ParameterKind
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
from umcglib.froc import *
from umcglib.binarize import dynamic_threshold


def print_p(*args, **kwargs):
    """
    Shorthand for print(..., flush=True)
    Useful on HPC cluster where output has buffered writes.
    """
    print(*args, **kwargs, flush=True)   


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


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

final_table = {}
difference = {}
for fold in range(5):

    if fold == 0:
        TEST_INDEX = [189,84]
    if fold == 1:
        TEST_INDEX = [828,12]
    if fold == 2:
        TEST_INDEX = [470,482]
    if fold == 3:
        TEST_INDEX = [591]
    if fold == 4:
        TEST_INDEX = [511,281,149]

    for img_idx in TEST_INDEX:

        for model in ['b800','b400']:
            
            image_paths = {}
            predictions_added = []
            segmentations_added = []
            images = []
            images_list = []
            segmentations = []
        
            if model is 'b800':
                MODEL_PATH = f'./../train_output/calc_exp_t2_b1400calc2_adccalc2_{fold}/models/calc_exp_t2_b1400calc2_adccalc2_{fold}.h5'
                # YAML_DIR = f'./../train_output/calc_exp_t2_b1400calc2_adccalc2_{fold}'
                # IMAGE_DIR = f'./../train_output/calc_exp_t2_b1400calc2_adccalc2_{fold}'
                SERIES = ['t2','b1400calc2','adccalc2']
            if model is 'b400':
                MODEL_PATH = f'./../train_output/calc_exp_t2_b1400calc3_adccalc3_{fold}/models/calc_exp_t2_b1400calc3_adccalc3_{fold}.h5'
                SERIES = ['t2','b1400calc3','adccalc3']

            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)

            pat_id = os.path.basename(os.path.normpath(seg_paths[img_idx]))[:-7] 

            # print_p("pat_idx:",pat_id)

            # print(image_paths['t2'][])
            # input('check?')
            # Read and preprocess each of the paths for each series, and the segmentations.
                # print('images number',[TEST_INDEX[img_idx]])
            img_s = {f'{s}': sitk.ReadImage(image_paths[s][img_idx], sitk.sitkFloat32) for s in SERIES}
            seg_s = sitk.ReadImage(seg_paths[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))


            ########### 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,axis=4))
            segmentations = move_dims(segmentations)
            # predictions = [preprocess_softmax(pred, threshold="dynamic")[0] for pred in predictions_blur]
            predictions = predictions_blur
            # print("the size of predictions is:",np.shape(predictions))
            # Remove outer edges 
            zeros = np.zeros(np.shape(predictions))
            test = predictions[:,2:-2,2:190,2:190]
            zeros[:,2:-2,2:190,2:190] = test
            predictions = zeros
            predictions = np.squeeze(predictions)
            # print('size pred:',np.shape(predictions))
            pred = sitk.GetImageFromArray(predictions)
            sitk.WriteImage(pred, f'../temp/lowest_pred_exp/worst/{pat_id}_{model}_pred.nii.gz')

            print('shape iamge:',np.shape(images_list))
            t2 = sitk.GetImageFromArray(np.squeeze(images_list)[:,:,:,0].T)
            sitk.WriteImage(t2, f'../temp/lowest_pred_exp/worst/{pat_id}_{model}_t2.nii.gz')

            b1400 = sitk.GetImageFromArray(np.squeeze(images_list)[:,:,:,1].T)
            sitk.WriteImage(b1400, f'../temp/lowest_pred_exp/worst/{pat_id}_{model}_b1400.nii.gz')

            adc = sitk.GetImageFromArray(np.squeeze(images_list)[:,:,:,2].T)
            sitk.WriteImage(adc, f'../temp/lowest_pred_exp/worst/{pat_id}_{model}_adc.nii.gz')