# Manuel A. Morales (moralesq@mit.edu)
# Harvard-MIT Department of Health Sciences & Technology  
# Athinoula A. Martinos Center for Biomedical Imaging

import numpy as np
from abc import ABC, abstractmethod
from tensorflow.keras.utils import Sequence
from scipy.ndimage.measurements import center_of_mass

import nibabel as nib
from dipy.align.reslice import reslice

class BaseDataset(Sequence, ABC):
    """This class is an abstract base class (ABC) for datasets."""

    def __init__(self, opt):
        self.opt  = opt
        self.root = opt.dataroot

    @abstractmethod
    def __len__(self):
        """Return the size of the dataset."""
        return 
    
    @abstractmethod
    def __getitem__(self, idx):
        """Return a data point and its metadata information."""
        pass
                       
class Transforms():
    
    def __init__(self, opt):
        self.opt = opt 
        self.transform, self.transform_inv = self.get_transforms(opt)
       
    def __crop__(self, x, inv=False):
        
        if inv:
            nx, ny = self.original_shape[:2]
            xinv = np.zeros(self.original_shape[:2] + x.shape[2:])
            xinv[nx//2-64:nx//2+64, ny//2-64:ny//2+64] += x
            return xinv
        else:
            nx, ny = x.shape[:2]
            return x[nx//2-64:nx//2+64, ny//2-64:ny//2+64]
    
    def __reshape_to_carson__(self, x, inv=False):
        
        if inv:
            if len(self.original_shape)==3:
                x = x.transpose(1,2,0,3)
            elif len(self.original_shape)==4:
                nx,ny,nz,nt=self.original_shape
                Nx, Ny = x.shape[1:3]
                x = x.reshape((nt, nz, Nx, Ny, self.opt.nlabels))
                x = x.transpose(2,3,1,0,4)                
        else:
            if len(x.shape) == 3:
                nx,ny,nz=x.shape
                x=x.transpose(2,0,1)
            elif len(x.shape) == 4:
                nx,ny,nz,nt=x.shape
                x=x.transpose(3,2,0,1)
                x=x.reshape((nt*nz,nx,ny))            
        return x

    def __reshape_to_carmen__(self, x, inv=False):
        if inv:
            x = np.concatenate((np.zeros(x[:1].shape), x))
            x = x.transpose((1,2,3,0,4)) 
        else:
            assert len(x.shape) == 4
            nx,ny,nz,nt=x.shape
            x=x.transpose(3,0,1,2)
            x=np.stack((np.repeat(x[:1],nt-1,axis=0), x[1:nt]), -1)
        return x  
    
    def __zscore__(self, x):

        if len(x.shape) == 3:
            axis=(1,2) # normalize in-plane images independently
        elif len(x.shape) == 5:
            axis=(1,2,3) # normalize volumes independently

        self.mu = x.mean(axis=axis, keepdims=True)
        self.sd = x.std(axis=axis, keepdims=True)
        return (x - self.mu)/(self.sd + 1e-8)

    def get_transforms(self, opt):

        transform_list     = []
        transform_inv_list = []
        if 'crop' in opt.preprocess:
            transform_list.append(self.__crop__)
            transform_inv_list.append(lambda x:self.__crop__(x,inv=True))
        if 'reshape_to_carson' in opt.preprocess:
            transform_list.append(self.__reshape_to_carson__)
            transform_inv_list.append(lambda x:self.__reshape_to_carson__(x,inv=True))
        elif 'reshape_to_carmen' in opt.preprocess:
            transform_list.append(self.__reshape_to_carmen__)
            transform_inv_list.append(lambda x:self.__reshape_to_carmen__(x,inv=True))
        if 'zscore' in opt.preprocess:
            transform_list.append(self.__zscore__)                
        
        return transform_list, transform_inv_list
          
    def apply(self, x):
        
        self.original_shape = x.shape
        for transform in self.transform:
            x = transform(x)
        return x
    
    def apply_inv(self, x):
        
        for transform in self.transform_inv[::-1]:
            x = transform(x)
        return x    
    

def _centercrop(x):
    nx, ny = x.shape[:2]
    return x[nx//2-64:nx//2+64,ny//2-64:ny//2+64]

def _roll(x,rx,ry):
    x = np.roll(x,rx,axis=0)
    x = np.roll(x,ry,axis=1)
    return x

def _roll2center(x, center):
    return _roll(x, int(x.shape[0]//2-center[0]), int(x.shape[1]//2-center[1]))
    
def _roll2center_crop(x, center):
    x = _roll2center(x, center)
    return _centercrop(x)
    
    
#####################################################
## FUNCTIONS TO ADD MORE FLEXIBILITY IN SEGMENTATION
#####################################################

def resample_nifti_inv(nifti_resampled, zooms, order=1, mode='nearest'):
    """ Resample `nifti_resampled` to `zooms` resolution.
    """
    data_resampled   = nifti_resampled.get_fdata()
    zooms_resampled  = nifti_resampled.header.get_zooms()[:3]
    affine_resampled = nifti_resampled.affine 
        
    data_resampled, affine_resampled = reslice(data_resampled, 
                                               affine_resampled, zooms_resampled, zooms, order=order, mode=mode)

    nifti = nib.Nifti1Image(data_resampled, affine_resampled)
    
    return nifti
    
def convert_back_to_nifti(data_resampled, nifti_info_subject, inv_256x256=False, order=1, mode='nearest'):

    if inv_256x256:
        data_resampled_mod_corr = roll_and_pad_256x256_to_center_inv(data_resampled, nifti_info=nifti_info_subject)
    else:
        data_resampled_mod_corr = data_resampled
        
    affine           = nifti_info_subject['affine']
    affine_resampled = nifti_info_subject['affine_resampled']
    zooms            = nifti_info_subject['zooms'][:3]
    zooms_resampled  = nifti_info_subject['zooms_resampled'][:3]
    
    data_resampled, affine_resampled = reslice(data_resampled_mod_corr, 
                                               affine_resampled, zooms_resampled, zooms, order=order, mode=mode)
    nifti = nib.Nifti1Image(data_resampled, affine_resampled)
    
    return nifti

def roll(x,rx,ry):
        x = np.roll(x,rx,axis=0)
        x = np.roll(x,ry,axis=1)
        return x
    
def roll2center(x, center):
    return roll(x, int(x.shape[0]//2-center[0]), int(x.shape[1]//2-center[1]))
    
def pad_256x256(x):
        xpad = (512-x.shape[0])//2, (512-x.shape[0])-(512-x.shape[0])//2
        ypad = (512-x.shape[1])//2, (512-x.shape[1])-(512-x.shape[1])//2
        pads = (xpad,ypad)+((0,0),)*(len(x.shape)-2)
        vals = ((0,0),)*len(x.shape)
        x = np.pad(x, pads, 'constant', constant_values=vals)
        x = x[512//2-256//2:512//2+256//2,512//2-256//2:512//2+256//2]
        return x
    
def roll_and_pad_256x256_to_center(x, center):
    x = roll2center(x, center)
    x = pad_256x256(x)
    return x

def roll_and_pad_256x256_to_center_inv(x, nifti_info):

    # Recover 256x256 array that was center-cropped to 128x128!
    x_256_256 = np.zeros((256,256)+x.shape[2:])
    x_256_256[128-64:128+64,128-64:128+64] += x
    
    # Coordinates to put the image in its original location.
    cx, cy         = nifti_info['center_resampled'][:2]
    cx_mod, cy_mod = nifti_info['center_resampled_256x256'][:2]
    
    x_inv = np.zeros(nifti_info['shape_resampled'][:3]+x.shape[3:])

    dx = min(int(cx),64)
    dy = min(int(cy),64)
    if (dx!=64)|(dy!=64):
        print('WARNING:FOV < 128x128!')

    x_inv[int(cx-dx):int(cx+dx),int(cy-dy):int(cy+dy)] += x_256_256[int(cx_mod-dx):int(cx_mod+dx),
                                                                    int(cy_mod-dy):int(cy_mod+dy)]
    return x_inv