import numpy as np 
import SimpleITK as sitk
import matplotlib.pyplot as plt
######## load images #############
# path_b50 = '/data/pca-rad/datasets/radboud_new/pat0351/2016/diffusie_cro/b-50/nifti_image.nii.gz'
# path_b400 = '/data/pca-rad/datasets/radboud_new/pat0351/2016/diffusie_cro/b-400/nifti_image.nii.gz'
# path_b800 = '/data/pca-rad/datasets/radboud_new/pat0351/2016/diffusie_cro/b-800/nifti_image.nii.gz'
# path_b1400 = '/data/pca-rad/datasets/radboud_new/pat0351/2016/diffusie_cro/b-1400/nifti_image.nii.gz'
# path_adc = '/data/pca-rad/datasets/radboud_new/pat0351/2016/dADC/nifti_image.nii.gz'

# path_b50 = 'X:/sfransen/train_output/adc_exp/b50.nii.gz'
# path_b400 = 'X:/sfransen/train_output/adc_exp/b400.nii.gz'
# path_b800 = 'X:/sfransen/train_output/adc_exp/b800.nii.gz'
# path_b1400 = 'X:/sfransen/train_output/adc_exp/b1400.nii.gz'
# path_adc = 'X:/sfransen/train_output/adc_exp/adc.nii.gz'

path_b50 = '/data/pca-rad/sfransen/train_output/adc_exp/b50_true.nii.gz'
path_b400 = '/data/pca-rad/sfransen/train_output/adc_exp/b400_true.nii.gz'
path_b800 = '/data/pca-rad/sfransen/train_output/adc_exp/b800_true.nii.gz'
path_b1400 = '/data/pca-rad/sfransen/train_output/adc_exp/b1400_true.nii.gz'
path_adc = '/data/pca-rad/sfransen/train_output/adc_exp/adc_calc_b50_b400_b800.nii.gz'

b50 = sitk.ReadImage(path_b50, sitk.sitkFloat32)
b50 = sitk.GetArrayFromImage(b50)
b400 = sitk.ReadImage(path_b400, sitk.sitkFloat32)
b400 = sitk.GetArrayFromImage(b400)
b800 = sitk.ReadImage(path_b800, sitk.sitkFloat32)
b800 = sitk.GetArrayFromImage(b800)
b1400 = sitk.ReadImage(path_b1400, sitk.sitkFloat32)
b1400 = sitk.GetArrayFromImage(b1400)
adc = sitk.ReadImage(path_adc, sitk.sitkFloat32)
adc = sitk.GetArrayFromImage(adc)

def show_img(greyscale_img):
    fig = plt.figure()
    plt.imshow(greyscale_img)
    plt.axis('on')
    path = f"iets.png"
    fig.savefig(path, dpi=300, bbox_inches='tight')

def calc_adc(b50, b400, b800):
    mean_dwi = (50 + 400 + 800) / 3
    mean_si = np.divide(np.add(np.add(np.log(b50), np.log(b400)), np.log(b800)), 3)

    denominator = np.multiply((50 - mean_dwi), np.subtract(np.log(b50), mean_si)) + np.multiply((400 - mean_dwi), np.subtract(np.log(b400), mean_si)) + np.multiply((800 - mean_dwi), np.subtract(np.log(b800), mean_si)) 
    numerator = np.power((50 - mean_dwi), 2) + np.power((400 - mean_dwi), 2) + np.power((800 - mean_dwi), 2)
    adc = np.divide(denominator, numerator)
    return adc * -1000000

def calc_adc_1(b50,b800):
    mean_dwi = (50 + 800) / 2
    mean_si = np.divide(np.add(np.log(b50), np.log(b800)), 2)

    denominator = np.multiply((50 - mean_dwi), np.subtract(np.log(b50), mean_si)) + np.multiply((800 - mean_dwi), np.subtract(np.log(b800), mean_si)) 
    numerator = np.power((50 - mean_dwi), 2) + np.power((800 - mean_dwi), 2)
    adc = np.divide(denominator, numerator)
    return adc * -1000000

def calc_adc_2(b50,b400):
    mean_dwi = (50 + 400) / 2
    mean_si = np.divide(np.add(np.log(b50), np.log(b400)), 2)

    denominator = np.multiply((50 - mean_dwi), np.subtract(np.log(b50), mean_si)) + np.multiply((400 - mean_dwi), np.subtract(np.log(b400), mean_si)) 
    numerator = np.power((50 - mean_dwi), 2) + np.power((400 - mean_dwi), 2)
    adc = np.divide(denominator, numerator)
    return adc * -1000000

def calc_adc_3(b400,b800):
    mean_dwi = (400 + 800) / 2
    mean_si = np.divide(np.add(np.log(b400), np.log(b800)), 2)

    denominator = np.multiply((400 - mean_dwi), np.subtract(np.log(b400), mean_si)) + np.multiply((800 - mean_dwi), np.subtract(np.log(b800), mean_si)) 
    numerator = np.power((400 - mean_dwi), 2) + np.power((800 - mean_dwi), 2)
    adc = np.divide(denominator, numerator)
    return adc * -1000000

def calc_high_b(b_value_high,b_value,b_image,ADC_map):
    high_b = np.multiply(b_image, np.log(np.multiply(np.subtract(b_value,b_value_high), ADC_map))) 
    return high_b


high_b = sitk.GetImageFromArray(b1400)
sitk.WriteImage(high_b, f"./../train_output/adc_exp/b1400_true.nii.gz")

high_b = calc_high_b(1400,50,b50,adc)
high_b = sitk.GetImageFromArray(high_b)
sitk.WriteImage(high_b, f"./../train_output/adc_exp/b1400_ref_b50.nii.gz")

high_b = calc_high_b(1400,400,b400,adc)
high_b = sitk.GetImageFromArray(high_b)
sitk.WriteImage(high_b, f"./../train_output/adc_exp/b1400_ref_b400.nii.gz")

high_b = calc_high_b(1400,800,b800,adc)
high_b = sitk.GetImageFromArray(high_b)
sitk.WriteImage(high_b, f"./../train_output/adc_exp/b1400_ref_b800.nii.gz")

# b50 = sitk.GetImageFromArray(b50)
# sitk.WriteImage(b50, "./../train_output/adc_exp/b50_true.nii.gz")

# b400 = sitk.GetImageFromArray(b400)
# sitk.WriteImage(b400, "./../train_output/adc_exp/b400_true.nii.gz")

# b800 = sitk.GetImageFromArray(b800)
# sitk.WriteImage(b800, "./../train_output/adc_exp/b800_true.nii.gz")

# b1400 = sitk.GetImageFromArray(b1400)
# sitk.WriteImage(b1400,f"./../train_output/adc_exp/b1400_true.nii.gz")


# adc = sitk.GetImageFromArray(adc)
# sitk.WriteImage(adc, f"adc_true.nii.gz")

# adc = calc_adc(b50,b400,b800)
# print("calculated with 3 adc shape:",adc.shape)
# adc = sitk.GetImageFromArray(adc)
# sitk.WriteImage(adc, f"adc_calc_b50_b400_b800.nii.gz")

# adc = calc_adc_1(b50,b800)
# print("calculated with 2 adc shape:",adc.shape)
# adc = sitk.GetImageFromArray(adc)
# sitk.WriteImage(adc, f"adc_calc_b50_b800.nii.gz")

# adc = calc_adc_2(b50,b400)
# print("calculated with 2 adc shape:",adc.shape)
# adc = sitk.GetImageFromArray(adc)
# sitk.WriteImage(adc, f"adc_calc_b50_b400.nii.gz")

# adc = calc_adc_3(b400,b800)
# print("calculated with 2 adc shape:",adc.shape)
# adc = sitk.GetImageFromArray(adc)
# sitk.WriteImage(adc, f"adc_calc_b400_b800.nii.gz")