import os
from os import path
from tqdm import tqdm
import numpy as np 
import SimpleITK as sitk

def calc_adc(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_with_zeros = np.divide(denominator, numerator) * -1000000
    adc = np.clip(adc_with_zeros,0,np.amax(adc_with_zeros))
    return adc

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

def append_new_line(file_name, text_to_append):
    """Append given text as a new line at the end of file"""
    # Open the file in append & read mode ('a+')
    with open(file_name, "a+") as file_object:
        # Move read cursor to the start of file.
        file_object.seek(0)
        # If file is not empty then append '\n'
        data = file_object.read(100)
        if len(data) > 0:
            file_object.write("\n")
        # Append text at the end of file
        file_object.write(text_to_append)



series = ['b50','b400']
# series = ['t2']

DATA_DIR = "./../data/Nijmegen paths/"

# Read the image paths from the data directory.
# Texts files are expected to have the name "[series_name].txt"
images, image_paths = {s: [] for s in series}, {}
segmentations = []

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)

for img_idx in tqdm(range(num_images)):
    img_b50 = sitk.ReadImage(image_paths['b50'][img_idx], sitk.sitkFloat32) 
    arr_b50 = sitk.GetArrayFromImage(img_b50)
    img_b400 = sitk.ReadImage(image_paths['b400'][img_idx], sitk.sitkFloat32) 
    arr_b400 = sitk.GetArrayFromImage(img_b400)
    # img_b800 = sitk.ReadImage(image_paths['b800'][img_idx], sitk.sitkFloat32) 
    # arr_b800 = sitk.GetArrayFromImage(img_b800)
    # img_t2 = sitk.ReadImage(image_paths['t2'][img_idx], sitk.sitkFloat32) 
    # img_t2 = sitk.GetArrayFromImage(img_t2)
    # img_seg = sitk.ReadImage(seg_paths[img_idx], sitk.sitkFloat32) 
    # img_seg = sitk.GetArrayFromImage(img_seg)

    adc = calc_adc(arr_b50,arr_b400)
    high_b = calc_high_b(1400,50,arr_b50,adc)
    
    adc = sitk.GetImageFromArray(adc)
    high_b = sitk.GetImageFromArray(high_b)
    # t2 = sitk.GetImageFromArray(img_t2)
    # seg = sitk.GetImageFromArray(img_seg)

    patient_ID = os.path.split(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(image_paths['b50'][img_idx])))))[1]
    path_adc = f'./../data/adc_calc_2/{patient_ID}.nii.gz'
    path_high_b = f'./../data/b1400_calc_2/{patient_ID}.nii.gz'
    # patient_ID = os.path.split(os.path.dirname(os.path.dirname(os.path.dirname(image_paths['t2'][img_idx]))))[1]
    # path_t2 = f'./../data/t2_calc/{patient_ID}.nii.gz'
    # path_seg = f'./../data/seg_calc/{patient_ID}.nii.gz'

    adc.CopyInformation(img_b50)
    high_b.CopyInformation(img_b50)
    sitk.WriteImage(adc, path_adc)
    sitk.WriteImage(high_b, path_high_b)
    # sitk.WriteImage(t2, path_t2)
    # sitk.WriteImage(seg, path_seg)

    append_new_line('adccalc2.txt',path_adc)
    append_new_line('b1400calc2.txt',path_high_b)
    # append_new_line('t2calc.txt',path_t2)
    # append_new_line('segcalc.txt',path_seg)