import numpy as np
import SimpleITK as sitk

from .helpers import *

def preprocess(imgs: dict, seg: sitk.Image, 
    shape: tuple, spacing: tuple, to_numpy=True) -> tuple:

    # Resample all of the images to the desired voxel spacing
    img_r = {k: resample(imgs[k], 
        min_shape=(s+1 for s in shape), 
        new_spacing=spacing) for k in imgs}
    seg_r = resample(seg, 
        min_shape=(s+1 for s in shape), 
        new_spacing=spacing,
        method=sitk.sitkNearestNeighbor)

    # Center crop one of the input images
    ref_seq = [k for k in img_r.keys()][0]
    ref_img = center_crop(img_r[ref_seq], shape=shape)

    # Then crop the remaining series / segmentation to match the input crop by
    # transforming them on the physical space of the cropped series.
    img_crop = {k: resample_to_reference(
        image=img_r[k], 
        ref_img=ref_img,
        interpolator=sitk.sitkLinear)
                    for k in img_r}
    seg_crop = resample_to_reference(
        image=seg_r, 
        ref_img=ref_img,
        interpolator=sitk.sitkNearestNeighbor)

    # Return sitk.Image instead of numpy np.ndarray.
    if not to_numpy:
        return img_crop, seg_crop
    
    img_n = {k: sitk.GetArrayFromImage(img_crop[k]).T for k in img_crop}
    seg_n = sitk.GetArrayFromImage(seg_crop).T
    seg_n = np.clip(seg_n, 0., 1.)

    # Z-score normalize all images
    # to do: 2*std
    for seq in img_n:
            img_n[seq] = (img_n[seq] - np.mean(img_n[seq])) / ( 2* np.std(img_n[seq]))
    return img_n, seg_n