"""
*Preliminary* pytorch implementation.

data generators for voxelmorph
"""

import numpy as np
import sys


def load_example_by_name(vol_name, seg_name=None):
    """
    load a specific volume and segmentation
    """
    X = np.load(vol_name)[ 'vol_data' ]
    X = np.reshape(X, (1,) + X.shape + (1,))

    return_vals = [X]

    if(seg_name):
        X_seg = np.load(seg_name)[ 'vol_data' ]
        X_seg = np.reshape(X_seg, (1,) + X_seg.shape + (1,))
        return_vals.append(X_seg)



    return tuple(return_vals)


def load_volfile(datafile):
    """
    load volume file
    formats: nii, nii.gz, mgz, npz
    if it's a npz (compressed numpy), assume variable names 'vol_data'
    """
    assert datafile.endswith(('.nii', '.nii.gz', '.mgz', '.npz')), 'Unknown data file: %s' % datafile

    if datafile.endswith(('.nii', '.nii.gz', '.mgz')):
        # import nibabel
        if 'nib' not in sys.modules:
            try:
                import nibabel as nib
            except:
                print('Failed to import nibabel. need nibabel library for these data file types.')

        X = nib.load(datafile).get_data()

    else:  # npz
        X = np.load(datafile)['vol_data']

    return X


def example_gen(vol_names, batch_size=1, return_segs=False, height=80):
    """
    generate examples

    Parameters:
        vol_names: a list or tuple of filenames
        batch_size: the size of the batch (default: 1)

        The following are fairly specific to our data structure, please change to your own (Apdated and Fixed)
        return_segs: logical on whether to return segmentations
    """

    while True:
        idxes = np.random.randint(len(vol_names), size=batch_size)

        X_data = []
        for idx in idxes:
            X = load_volfile(vol_names[idx])
            X = X[np.newaxis, ..., np.newaxis]
            X_data.append(X[:,:,:,height,:])

        if batch_size > 1:
            return_vals = [np.concatenate(X_data, 0)]
        else:
            return_vals = [X_data[0]]

        # also return segmentations
        if return_segs:
            X_data = []
            for idx in idxes:
                v = vol_names[idx].replace('norm', 'aseg')
                v = v.replace('vols', 'asegs')
                X_seg = load_volfile(v)
                X_seg = X_seg[np.newaxis, ..., np.newaxis]
                #We should take into consideration the height dim , even in the segementation
                X_data.append(X_seg[:,:,:,height,:]) 

            if batch_size > 1:
                return_vals.append(np.concatenate(X_data, 0))
            else:
                return_vals.append(X_data[0])

        yield tuple(return_vals)