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
from typing import Dict, Union
import SimpleITK as sitk
import numpy as np

import umcglib.images as im

def calculate_adc(b_dict: Dict[int, Union[sitk.Image, np.ndarray]]):
    "Calculates the adc based on b50, b400 and b800 DWI images."
    b_values = [int(b) for b in b_dict]

    # Determine the type of inputs (Image or ndarray)
    b_images = list(b_dict.values())
    sample = b_images[0]

    if type(sample) == sitk.Image:
        # Extract numpy arrays, and call function again with numpy input
        b_dict_n = {b: sitk.GetArrayFromImage(b_dict[b]).T for b in b_dict}
        adc_n = calculate_adc(b_dict_n)

        # Convert calculated map back to sitk
        return im.to_sitk(adc_n, ref_img=sample)

    mean_b_value = sum(b_values) / len(b_values)
    mean_si = sum(b_images) / len(b_values)

    denominator = np.zeros_like(sample)
    numerator = 0
    for b in b_dict:
        b_img = b_dict[b]
        denominator += (b - mean_b_value) * (np.log(b_img) - mean_si)
        numerator += (b - mean_b_value) ** 2

    adc_with_zeros = denominator / numerator * -1000000
    adc = np.clip(adc_with_zeros, a_min=0, a_max=None)
    return adc

img_b0 = '../../datasets/anonymized_mri/only_nii_directory/108-M-108/2018/Diffusie_b0-50-400-800-2000/b-0/701.nii.gz'
img_b50 = '../../datasets/anonymized_mri/only_nii_directory/108-M-108/2018/Diffusie_b0-50-400-800-2000/b-50/701.nii.gz'
img_b400 = '../../datasets/anonymized_mri/only_nii_directory/108-M-108/2018/Diffusie_b0-50-400-800-2000/b-400/701.nii.gz'
img_b800 = '../../datasets/anonymized_mri/only_nii_directory/108-M-108/2018/Diffusie_b0-50-400-800-2000/b-800/701.nii.gz'
img_b2000 = '../../datasets/anonymized_mri/only_nii_directory/108-M-108/2018/Diffusie_b0-50-400-800-2000/b-2000/701.nii.gz'
adc_or = '../../datasets/anonymized_mri/only_nii_directory/108-M-108/2018/dADC/702_.nii.gz'

img_b0 = sitk.ReadImage(img_b0, sitk.sitkFloat32)
img_b0 = sitk.GetArrayFromImage(img_b0)

img_b50 = sitk.ReadImage(img_b50, sitk.sitkFloat32)
img_b50 = sitk.GetArrayFromImage(img_b50)

img_b400 = sitk.ReadImage(img_b400, sitk.sitkFloat32)
img_b400 = sitk.GetArrayFromImage(img_b400)

img_b800 = sitk.ReadImage(img_b800, sitk.sitkFloat32)
img_b800 = sitk.GetArrayFromImage(img_b800)

img_b2000 = sitk.ReadImage(img_b2000, sitk.sitkFloat32)
img_b2000 = sitk.GetArrayFromImage(img_b2000)

ADC_original_img = sitk.ReadImage(adc_or, sitk.sitkFloat32)
ADC_original = sitk.GetArrayFromImage(ADC_original_img)

print(np.shape(ADC_original))
print(np.shape(img_b800))

ADC_b0_b800 = calculate_adc({0:img_b0,800:img_b800})
ADC_b50_b800 = calculate_adc({50:img_b50,800:img_b800})
ADC_b50_b400_b800 = calculate_adc({50:img_b50,400:img_b400,800:img_b800})
ADC_b0_b50_b400_b800 = calculate_adc({0:img_b0,50:img_b50,400:img_b400,800:img_b800})

ADC_b50_2000 = calculate_adc({0:img_b0,800:img_b800,2000:img_b2000})
ADC_b0_b800_2000 = calculate_adc({0:img_b0,800:img_b800,2000:img_b2000})
ADC_b50_b800_2000 = calculate_adc({50:img_b50,800:img_b800,2000:img_b2000})
ADC_b50_b400_b800_2000 = calculate_adc({50:img_b50,400:img_b400,800:img_b800,2000:img_b2000})
ADC_b0_b50_b400_b800_2000 = calculate_adc({0:img_b0,50:img_b50,400:img_b400,800:img_b800,2000:img_b2000})


ADC_b50_b400_b800_2000 = sitk.GetImageFromArray(ADC_b50_b400_b800_2000)
ADC_b50_b400_b800_2000.CopyInformation(ADC_original_img)
sitk.WriteImage(ADC_b50_b400_b800_2000, f"../train_output/ADC_b50_b400_b800_2000.nii.gz")

SLIDE = 10
diffs = []
diffs.append(ADC_b0_b800[SLIDE,:,:]-ADC_original[SLIDE,:,:])
diffs.append(ADC_b50_b800[SLIDE,:,:]-ADC_original[SLIDE,:,:])
diffs.append(ADC_b50_b400_b800[SLIDE,:,:]-ADC_original[SLIDE,:,:])
diffs.append(ADC_b0_b50_b400_b800[SLIDE,:,:]-ADC_original[SLIDE,:,:])

diffs_2000 = []
diffs_2000.append(ADC_b0_b800_2000[SLIDE,:,:]-ADC_original[SLIDE,:,:])
diffs_2000.append(ADC_b50_b800_2000[SLIDE,:,:]-ADC_original[SLIDE,:,:])
diffs_2000.append(ADC_b50_b400_b800_2000[SLIDE,:,:]-ADC_original[SLIDE,:,:])
diffs_2000.append(ADC_b0_b50_b400_b800_2000[SLIDE,:,:]-ADC_original[SLIDE,:,:])


TITLES = ['ADC_b0_b800','ADC_b50_b800','ADC_b50_b400_b800','ADC_b0_b50_b400_b800']
TITLES_2000 = ['ADC_b0_b800_2000','ADC_b50_b800_2000','ADC_b50_b400_b800_2000','ADC_b0_b50_b400_b800_2000']

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

x = abs(np.asarray(np.stack(diffs_2000)))
x[np.isneginf(x)] = 0
x[np.isposinf(x)] = 0

vmin = np.nanmin(x)
vmax = np.nanmax(x)
print('vmax',vmax,'  ',type(np.asarray(np.stack(diffs))))
for indx in range(len(TITLES)):
    print(indx)
    axes[indx,0].imshow(abs(np.transpose(diffs[indx])),cmap='gray',vmin =vmin , vmax=vmax)
    axes[indx,0].set_title(f'{TITLES[indx]}')
    axes[indx,0].set_axis_off()
    axes[indx,0].set_axis_off()
    axes[indx,1].imshow(abs(np.transpose(diffs_2000[indx])),cmap='gray',vmin =vmin , vmax=vmax)
    axes[indx,1].set_title(f'{TITLES_2000[indx]}')
    axes[indx,1].set_axis_off()
    axes[indx,1].set_axis_off()

plt.savefig(f'./../train_output/adc_diff_joeran.png', dpi=300)