from pathlib import Path
import argparse
import numpy as np
import nibabel as nib
from PIL import Image, ImageSequence
from scipy import ndimage
from skimage import morphology, filters
from nipype.interfaces.ants.segmentation import N4BiasFieldCorrection

import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)


def get_out_paths(out_dir, stem):
    out_path_scaled = Path(out_dir, f'{stem}.nii')
    out_path_unscaled = Path(out_dir, f'{stem}_us.nii')
    return out_path_scaled, out_path_unscaled


def bias_field_correction(image, output, dim=2):
    print('Running N4BiasFieldCorrection on input slice.')
    N4BiasFieldCorrection(
        input_image=image,
        output_image=output,
        dimension=dim
    ).run()


def image_data_to_nii(pixdim, image_data, shrink, out_dir, file_path, save_unscaled=False):
    image_dim = (image_data.shape[0], image_data.shape[1])
    scale = 1/shrink
    new_dim = (round(image_dim[1] * scale), round(image_dim[0] * scale))
    new_arr = np.ndarray(new_dim + (image_data.shape[2],))
    print(f'Pixel dimensions: {pixdim} mm')

    for i in range(0,image_data.shape[2]):
        cur_channel = image_data[:, :, i]
        resized = np.array(Image.fromarray(cur_channel).resize(new_dim)).transpose()
        new_arr[:, :, i] = resized
    path_scaled, path_unscaled = get_out_paths(out_dir, Path(file_path).stem)
    nii_scaled = nib.Nifti1Image(new_arr, np.eye(4))
    nii_scaled.header['pixdim'][1:3] = pixdim * shrink, pixdim * shrink
    nib.save(nii_scaled, str(path_scaled))

    if save_unscaled:
        nii_unscaled = nib.Nifti1Image(image_data, np.eye(4))
        nii_unscaled.header['xyzt_units'] = 3
        nii_unscaled.header['pixdim'][1:3] = pixdim, pixdim
        nib.save(nii_unscaled, str(path_unscaled))
    print(f'Preprocessed:  {path_scaled}\n')
    return path_scaled, path_unscaled


def tiff_to_nii(tif_path, out_dir, pixdim=None, shrink=10):
    Image.MAX_IMAGE_PIXELS = None
    tif_image = Image.open(tif_path)
    tif_header = dict(tif_image.tag)
    output = np.empty(np.array(tif_image).shape + (0,))
    if not pixdim:
        pixdim = 10e2/tif_header[282][0][0]
    for i, page in enumerate(ImageSequence.Iterator(tif_image)):
        page_data = np.expand_dims(np.array(page), 2)
        output = np.concatenate((output, page_data), 2)
    return image_data_to_nii(pixdim, output, shrink, out_dir, tif_path)


def split_nii_channels(nii_path, out_dir=None, flip=False, mask_index=-1, bias=True):
    if out_dir is None:
        out_dir = nii_path.parent
    nii = nib.load(str(nii_path))
    nii_data = nii.get_fdata()
    nii_header = nii.header

    if mask_index == -1:
        mask_index = nii_data.shape[2] - 1
    paths = []

    for i in range(0, nii_data.shape[2]):
        out_path = out_dir / f'im_c{i+1}.nii'
        channel_data = nii_data[:, :, i]

        if flip:
            channel_data = np.flip(channel_data, 1)
        if i == mask_index:
            channel_data = mask_foreground(channel_data)

        new_header = nii_header
        new_header['dim'][0] = 2
        nii = nib.Nifti1Image(channel_data, np.eye(4), header=new_header)
        nib.save(nii, str(out_path))

        if i == mask_index and bias:
            bias_field_correction(str(out_path), str(out_path))
            corrected = nib.load(str(out_path))
            corrected_data = corrected.get_fdata()
            corrected_normalized = corrected_data / np.mean(corrected_data[corrected_data != 0])
            nii_corrected = nib.Nifti1Image(corrected_normalized, corrected.affine, corrected.header)
            nib.save(nii_corrected, str(out_path))
        paths.append(out_path)
    return paths


def mask_foreground(raw_data):
    raw_max = raw_data.max()
    raw_data = raw_data / raw_max
    blurred_data = ndimage.gaussian_filter(raw_data, 4)

    threshold = filters.threshold_otsu(raw_data) / 2
    threshold_data = blurred_data > threshold
    connected_structure = ndimage.generate_binary_structure(2, 2)  # Connects adjacent and diagonal.
    padded_comp, padded_nr = ndimage.label(threshold_data, structure=connected_structure)

    comps, comps_count = np.unique(padded_comp, return_counts=True)
    comps_count, comps = zip(*sorted(zip(comps_count, comps), reverse=True))

    two_biggest_cc = ((comps[0], np.average(comps[0])), (comps[1], np.average(comps[1])))
    biggest_cc = max(two_biggest_cc, key=lambda a: a[1])[0]

    foreground_mask = np.where(padded_comp == biggest_cc, True, False)
    closed = morphology.binary_closing(foreground_mask, selem=morphology.square(30))
    raw_data = np.where(closed, raw_data, 0)
    return raw_data * raw_max
    
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("file", help="Location of file to process")
    parser.add_argument("dir", help="Directory for preprocessed Files")
    parser.add_argument("-s", "--series", type=int, help="Series to extract")
    parser.add_argument('-b', type=bool, default= True, help='Bias field correct image')
    parser.add_argument('--pdim', type=float, help='Pixel dimensions (Retrieved from Tiff file if not set)')
    args = parser.parse_args()

    filetype = Path(args.file).suffix.lower()
    if filetype == '.tif' or filetype == '.tiff':
        n_path = tiff_to_nii(Path(args.file), Path(args.dir))
    elif filetype == '.nd2':
        n_path = nd2_to_nii(Path(args.file), Path(args.dir), args.series)
    elif filetype == '.nii':
        n_path = Path(args.file)