import os 
from os import path
import SimpleITK as sitk
from sfransen.DWI_exp import preprocess
import numpy as np 
import matplotlib.pyplot as plt

SERIES = ['adccalc2','adccalc','adccalc3']
TARGET_SPACING = (0.5, 0.5, 3) 
INPUT_SHAPE = (192, 192, 24, len(SERIES))
IMAGE_SHAPE = INPUT_SHAPE[:3]

DATA_DIR = "./../data/Nijmegen paths/"
########## 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)

img_idx = 371

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

# 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.squeeze(np.transpose(images_list, (0, 2, 3, 4, 1)))

print(np.shape(images_list))
SLIDE = 10
diff1 = images_list[:,:,SLIDE,0]-images_list[:,:,SLIDE,1]
diff2 = images_list[:,:,SLIDE,0]-images_list[:,:,SLIDE,2]
diff3 = images_list[:,:,SLIDE,1]-images_list[:,:,SLIDE,2]

fig, axes = plt.subplots(2,len(SERIES))
fig.suptitle('ADC map differences', fontsize=16)

maxx_2 = np.amax(images_list[:,:,SLIDE,:])
minn_2 = np.amin(images_list[:,:,SLIDE,:])

TITLES = ['$1) ADC_{b50-b400}$','$2) ADC_{b50-b800}$','$3) ADC_{b50-b400-b800}$']
for indx in range(len(SERIES)):
    print(indx)
    im = axes[0,indx].imshow(np.transpose(images_list[:,:,SLIDE,indx]),cmap='gray',vmin=minn_2,vmax=maxx_2)
    axes[0,indx].set_title(TITLES[indx])
    axes[0,indx].set_axis_off()
    axes[0,indx].set_axis_off()

maxx_2 = np.amax(images_list[:,:,SLIDE,:])
minn_2 = np.amin(images_list[:,:,SLIDE,:])
# print('>>>>>>',minn_2)
cbar_ax = fig.add_axes([0.85, 0.57, 0.02, 0.25])
cbar2 = fig.colorbar(im, cax=cbar_ax, ticks=[minn_2+0.001, maxx_2-0.001])
cbar2.ax.set_yticklabels([f'0', f'{1.3}'])


maxx = np.amax([np.amax(diff1), np.amax(diff2), np.amax(diff3)]) *0.99
minn = np.amin([np.amin(diff1), np.amin(diff2), np.amin(diff3)])

im = axes[1,0].imshow(np.transpose(diff1),cmap='gray',vmax = maxx,vmin=minn)
axes[1,0].set_axis_off()
axes[1,0].set_axis_off()
axes[1,0].set_title('Protocol 1-2')
im = axes[1,1].imshow(np.transpose(diff2),cmap='gray',vmax = maxx,vmin=minn)
axes[1,1].set_axis_off()
axes[1,1].set_axis_off()
axes[1,1].set_title('Protocol 1-3')

im = axes[1,2].imshow(np.transpose(diff3),cmap='gray',vmax = maxx,vmin=minn)
axes[1,2].set_axis_off()
axes[1,2].set_axis_off()
axes[1,2].set_title('Protocol 2-3')

fig.subplots_adjust(right=0.8)
cbar_ax = fig.add_axes([0.85, 0.15, 0.02, 0.25])
cbar = fig.colorbar(im, cax=cbar_ax, ticks=[minn, 0, maxx])
cbar.ax.set_yticklabels([f'{round(minn,1)}','0', f'{round(maxx,1)}'])
plt.savefig(f'./../train_output/adc_diff.png', dpi=300)