# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
"""
Intensity-based image registration
"""
from __future__ import absolute_import
from __future__ import print_function

import os
import numpy as np
import scipy.ndimage as nd
from nibabel import Nifti1Image

from .optimizer import configure_optimizer
from .affine import inverse_affine, subgrid_affine, affine_transforms
from .chain_transform import ChainTransform
from .similarity_measures import similarity_measures as builtin_simi
from ._register import _joint_histogram

MAX_INT = np.iinfo(np.intp).max

# Module globals
VERBOSE = os.environ.get('NIREG_DEBUG_PRINT', False)  # enables online print statements
CLAMP_DTYPE = 'short'  # do not edit
NPOINTS = 64 ** 3

# Dictionary of interpolation methods (partial volume, trilinear,
# random)
interp_methods = {'pv': 0, 'tri': 1, 'rand': -1}


def unpack(val, numtype):
    try:
        tmp = numtype(val)
        out = (tmp, tmp)
    except:
        out = (numtype(val[0]), numtype(val[1]))
    return out


class HistogramRegistration(object):
    """
    A class to reprensent a generic intensity-based image registration
    algorithm.
    """
    def __init__(self, from_img, to_img,
                 from_mask=None,
                 to_mask=None,
                 bins=256,
                 spacing=None,
                 similarity='crl1',
                 interp='pv',
                 sigma=0,
                 renormalize=False,
                 dist=None):
        """Creates a new histogram registration object.

        Parameters
        ----------
        from_img : nibabel image
          `From` image
        to_img : nibabel image
          `To` image
        from_mask : array-like
          Mask to apply to the `from` image
        to_mask : array-like
          Mask to apply to the `to` image
        bins : integer or sequence
          Number of histogram bins to represent the `from` and `to`
          image, respectively. If float, the same binning is applied
          to both images.
        spacing : None or sequence
          A sequence of three integers representing the subsampling
          factors applied to the `from` image grid for faster
          similarity computation. If None, the spacing is set
          automatically so as to trade off between registration
          accuracy and computation time.
        similarity : str or callable
          Cost-function for assessing image similarity. If a string,
          one of 'cc': correlation coefficient, 'cr': correlation
          ratio, 'crl1': L1-norm based correlation ratio, 'mi': mutual
          information, 'nmi': normalized mutual information, 'slr':
          supervised log-likelihood ratio. If a callable, it should
          take a two-dimensional array representing the image joint
          histogram as an input and return a float.
       dist: None or array-like
          Joint intensity probability distribution model for use with the 
          'slr' measure. Should be of shape (from_bins, to_bins).
       interp : str
         Interpolation method.  One of 'pv': Partial volume, 'tri':
         Trilinear, 'rand': Random interpolation.  See ``joint_histogram.c``
       sigma : float or sequence
         Standard deviation(s) in millimeters of isotropic Gaussian
         kernels used to smooth the `from` and `to` images,
         respectively. If float, the same kernel size is applied to
         both images. If 0, no smoothing is applied.
        """
        # Binning sizes
        from_bins, to_bins = unpack(bins, int)

        # Smoothing kernel sizes
        self._from_sigma, self._to_sigma = unpack(sigma, float)

        # Clamping of the `from` image. The number of bins may be
        # overriden if unnecessarily large.
        data, from_bins_adjusted = clamp(from_img,
                                         from_bins,
                                         mask=from_mask,
                                         sigma=self._from_sigma)
        if not similarity == 'slr':
            from_bins = from_bins_adjusted
        self._from_img = Nifti1Image(data, from_img.get_affine())

        # Set field of view in the `from` image with potential
        # subsampling for faster similarity evaluation. This also sets
        # the _from_data and _vox_coords attributes
        if spacing == None:
            npoints = NPOINTS
        else:
            npoints = None
        if from_mask == None:
            corner, size = (0, 0, 0), None
        else:
            corner, size = smallest_bounding_box(from_mask)
        self.set_fov(spacing=spacing, corner=corner, size=size, 
                     npoints=npoints)

        # Clamping of the `to` image including padding with -1
        data, to_bins_adjusted = clamp(to_img,
                                       to_bins,
                                       mask=to_mask,
                                       sigma=self._to_sigma)
        if not similarity == 'slr':
            to_bins = to_bins_adjusted
        self._to_data = -np.ones(np.array(to_img.shape) + 2, dtype=CLAMP_DTYPE)
        self._to_data[1:-1, 1:-1, 1:-1] = data
        self._to_inv_affine = inverse_affine(to_img.get_affine())

        # Joint histogram: must be double contiguous as it will be
        # passed to C routines which assume so
        self._joint_hist = np.zeros([from_bins, to_bins], dtype='double')

        # Set default registration parameters
        self._set_interp(interp)
        self._set_similarity(similarity, renormalize, dist=dist)

    def _get_interp(self):
        return list(interp_methods.keys())[\
            list(interp_methods.values()).index(self._interp)]

    def _set_interp(self, interp):
        self._interp = interp_methods[interp]

    interp = property(_get_interp, _set_interp)

    def set_fov(self, spacing=None, corner=(0, 0, 0), size=None,
                npoints=None):
        """
        Defines a subset of the `from` image to restrict joint
        histogram computation.

        Parameters
        ----------
        spacing : sequence (3,) of positive integers
          Subsampling of image in voxels, where None (default) results
          in the subsampling to be automatically adjusted to roughly
          match a cubic grid with `npoints` voxels
        corner : sequence (3,) of positive integers
          Bounding box origin in voxel coordinates
        size : sequence (3,) of positive integers
          Desired bounding box size
        npoints : positive integer
          Desired number of voxels in the bounding box. If a `spacing`
          argument is provided, then `npoints` is ignored.
        """
        if spacing is None and npoints is None:
            spacing = [1, 1, 1]
        if size is None:
            size = self._from_img.shape
        slicer = lambda c, s, sp:\
            tuple([slice(c[i], s[i] + c[i], sp[i]) for i in range(3)])
        # Adjust spacing to match desired field of view size
        if spacing is not None:
            fov_data = self._from_img.get_data()[slicer(corner, size, spacing)]
        else:
            fov_data = self._from_img.get_data()[slicer(corner, size, [1, 1, 1])]
            spacing = ideal_spacing(fov_data, npoints=npoints)
            fov_data = self._from_img.get_data()[slicer(corner, size, spacing)]
        self._from_data = fov_data
        self._from_npoints = (fov_data >= 0).sum()
        self._from_affine = subgrid_affine(self._from_img.get_affine(),
                                           slicer(corner, size, spacing))
        # We cache the voxel coordinates of the clamped image
        self._from_spacing = spacing
        self._vox_coords =\
            np.indices(self._from_data.shape).transpose((1, 2, 3, 0))

    def _set_similarity(self, similarity, renormalize=False, dist=None):
        if similarity in builtin_simi:
            if similarity == 'slr':
                if dist is None:
                    raise ValueError('slr measure requires a joint intensity distribution model, '
                                     'see `dist` argument of HistogramRegistration')
                if dist.shape != self._joint_hist.shape:
                    raise ValueError('Wrong shape for the `dist` argument')
            self._similarity = similarity
            self._similarity_call =\
                builtin_simi[similarity](self._joint_hist.shape,
                                         self._from_npoints,
                                         renormalize=renormalize,
                                         dist=dist)
        else:
            if not hasattr(similarity, '__call__'):
                raise ValueError('similarity should be callable')
            self._similarity = 'custom'
            self._similarity_call = similarity

    def _get_similarity(self):
        return self._similarity

    similarity = property(_get_similarity, _set_similarity)

    def eval(self, T):
        """
        Evaluate similarity function given a world-to-world transform.

        Parameters
        ----------
        T : Transform
            Transform object implementing ``apply`` method
        """
        Tv = ChainTransform(T, pre=self._from_affine, post=self._to_inv_affine)
        return self._eval(Tv)

    def eval_gradient(self, T, epsilon=1e-1):
        """
        Evaluate the gradient of the similarity function wrt
        transformation parameters.

        The gradient is approximated using central finite differences
        at the transformation specified by `T`. The input
        transformation object `T` is modified in place unless it has a
        ``copy`` method.

        Parameters
        ----------
        T : Transform
            Transform object implementing ``apply`` method
        epsilon : float
            Step size for finite differences in units of the
            transformation parameters

        Returns
        -------
        g : ndarray
            Similarity gradient estimate
        """
        param0 = T.param.copy()
        if hasattr(T, 'copy'):
            T = T.copy()

        def simi(param):
            T.param = param
            return self.eval(T)

        return approx_gradient(simi, param0, epsilon)

    def eval_hessian(self, T, epsilon=1e-1, diag=False):
        """
        Evaluate the Hessian of the similarity function wrt
        transformation parameters.

        The Hessian or its diagonal is approximated at the
        transformation specified by `T` using central finite
        differences. The input transformation object `T` is modified
        in place unless it has a ``copy`` method.

        Parameters
        ----------
        T : Transform
            Transform object implementing ``apply`` method
        epsilon : float
            Step size for finite differences in units of the
            transformation parameters
        diag : bool
            If True, approximate the Hessian by a diagonal matrix.

        Returns
        -------
        H : ndarray
            Similarity Hessian matrix estimate
        """
        param0 = T.param.copy()
        if hasattr(T, 'copy'):
            T = T.copy()

        def simi(param):
            T.param = param
            return self.eval(T)

        if diag:
            return np.diag(approx_hessian_diag(simi, param0, epsilon))
        else:
            return approx_hessian(simi, param0, epsilon)

    def _eval(self, Tv):
        """
        Evaluate similarity function given a voxel-to-voxel transform.

        Parameters
        ----------
        Tv : Transform
             Transform object implementing ``apply`` method
             Should map voxel space to voxel space
        """
        # trans_vox_coords needs be C-contiguous
        trans_vox_coords = Tv.apply(self._vox_coords)
        interp = self._interp
        if self._interp < 0:
            interp = - np.random.randint(MAX_INT)
        _joint_histogram(self._joint_hist,
                         self._from_data.flat,  # array iterator
                         self._to_data,
                         trans_vox_coords,
                         interp)
        # Make sure all joint histogram entries are non-negative
        np.maximum(self._joint_hist, 0, self._joint_hist)
        return self._similarity_call(self._joint_hist)

    def optimize(self, T, optimizer='powell', xtol=1e-2, ftol=1e-2, gtol=1e-3,
                 maxiter=25, maxfun=None, **kwargs):
        """ Optimize transform `T` with respect to similarity measure.

        The input object `T` will change as a result of the optimization.

        Parameters
        ----------
        T : object or str
          An object representing a transformation that should
          implement ``apply`` method and ``param`` attribute or
          property. If a string, one of 'rigid', 'similarity', or
          'affine'. The corresponding transformation class is then
          initialized by default.
        optimizer : str
          Name of optimization function (one of 'powell', 'steepest',
          'cg', 'bfgs', 'simplex')
        **kwargs : dict
          keyword arguments to pass to optimizer

        Returns
        -------
        T : object
          Locally optimal transformation
        """
        # Replace T if a string is passed
        if T in affine_transforms:
            T = affine_transforms[T]()

        # Pull callback out of keyword arguments, if present
        callback = kwargs.pop('callback', None)

        # Create transform chain object with T generating params
        Tv = ChainTransform(T, pre=self._from_affine, post=self._to_inv_affine)
        tc0 = Tv.param

        # Cost function to minimize
        def cost(tc):
            # This is where the similarity function is calculcated
            Tv.param = tc
            return -self._eval(Tv)

        # Callback during optimization
        if callback is None and VERBOSE:

            def callback(tc):
                Tv.param = tc
                print(Tv.optimizable)
                print(str(self.similarity) + ' = %s' % self._eval(Tv))
                print('')

        # Switching to the appropriate optimizer
        if VERBOSE:
            print('Initial guess...')
            print(Tv.optimizable)

        kwargs.setdefault('xtol', xtol)
        kwargs.setdefault('ftol', ftol)
        kwargs.setdefault('gtol', gtol)
        kwargs.setdefault('maxiter', maxiter)
        kwargs.setdefault('maxfun', maxfun)

        fmin, args, kwargs = configure_optimizer(optimizer,
                                                 fprime=None,
                                                 fhess=None,
                                                 **kwargs)

        # Output
        if VERBOSE:
            print('Optimizing using %s' % fmin.__name__)
        kwargs['callback'] = callback
        Tv.param = fmin(cost, tc0, *args, **kwargs)
        return Tv.optimizable

    def explore(self, T, *args):
        """
        Evaluate the similarity at the transformations specified by
        sequences of parameter values.

        For instance:

        s, p = explore(T, (0, [-1,0,1]), (4, [-2.,2]))

        Parameters
        ----------
        T : object
          Transformation around which the similarity function is to be
          evaluated. It is modified in place unless it has a ``copy``
          method.
        args : tuple
          Each element of `args` is a sequence of two elements, where
          the first element specifies a transformation parameter axis
          and the second element gives the successive parameter values
          to evaluate along that axis.

        Returns
        -------
        s : ndarray
          Array of similarity values
        p : ndarray
          Corresponding array of evaluated transformation parameters
        """
        nparams = T.param.size
        if hasattr(T, 'copy'):
            T = T.copy()
        deltas = [[0] for i in range(nparams)]
        for a in args:
            deltas[a[0]] = a[1]
        grids = np.mgrid[[slice(0, len(d)) for d in deltas]]
        ntrials = np.prod(grids.shape[1:])
        Deltas = [np.asarray(deltas[i])[grids[i, :]].ravel()\
                      for i in range(nparams)]
        simis = np.zeros(ntrials)
        params = np.zeros([nparams, ntrials])

        Tv = ChainTransform(T, pre=self._from_affine,
                            post=self._to_inv_affine)
        param0 = Tv.param
        for i in range(ntrials):
            param = param0 + np.array([D[i] for D in Deltas])
            Tv.param = param
            simis[i] = self._eval(Tv)
            params[:, i] = param

        return simis, params


def ideal_spacing(data, npoints):
    """
    Tune spacing factors so that the number of voxels in the
    output block matches a given number.

    Parameters
    ----------
    data : ndarray or sequence
      Data image to subsample
    npoints : number
      Target number of voxels (negative values will be ignored)

    Returns
    -------
    spacing: ndarray
      Spacing factors
    """
    dims = data.shape
    actual_npoints = (data >= 0).sum()
    spacing = np.ones(3, dtype='uint')

    while actual_npoints > npoints:

        # Subsample the direction with the highest number of samples
        ddims = dims / spacing
        if ddims[0] >= ddims[1] and ddims[0] >= ddims[2]:
            dir = 0
        elif ddims[1] > ddims[0] and ddims[1] >= ddims[2]:
            dir = 1
        else:
            dir = 2
        spacing[dir] += 1
        subdata = data[::spacing[0], ::spacing[1], ::spacing[2]]
        actual_npoints = (subdata >= 0).sum()

    return spacing


def smallest_bounding_box(msk):
    """
    Extract the smallest bounding box from a mask

    Parameters
    ----------
    msk : ndarray
      Array of boolean

    Returns
    -------
    corner: ndarray
      3-dimensional coordinates of bounding box corner
    size: ndarray
      3-dimensional size of bounding box
    """
    x, y, z = np.where(msk > 0)
    corner = np.array([x.min(), y.min(), z.min()])
    size = np.array([x.max() + 1, y.max() + 1, z.max() + 1])
    return corner, size


def approx_gradient(f, x, epsilon):
    """
    Approximate the gradient of a function using central finite
    differences

    Parameters
    ----------
    f: callable
      The function to differentiate
    x: ndarray
      Point where the function gradient is to be evaluated
    epsilon: float
      Stepsize for finite differences

    Returns
    -------
    g: ndarray
      Function gradient at `x`
    """
    n = len(x)
    g = np.zeros(n)
    ei = np.zeros(n)
    for i in range(n):
        ei[i] = .5 * epsilon
        g[i] = (f(x + ei) - f(x - ei)) / epsilon
        ei[i] = 0
    return g


def approx_hessian_diag(f, x, epsilon):
    """
    Approximate the Hessian diagonal of a function using central
    finite differences

    Parameters
    ----------
    f: callable
      The function to differentiate
    x: ndarray
      Point where the Hessian is to be evaluated
    epsilon: float
      Stepsize for finite differences

    Returns
    -------
    h: ndarray
      Diagonal of the Hessian at `x`
    """
    n = len(x)
    h = np.zeros(n)
    ei = np.zeros(n)
    fx = f(x)
    for i in range(n):
        ei[i] = epsilon
        h[i] = (f(x + ei) + f(x - ei) - 2 * fx) / (epsilon ** 2)
        ei[i] = 0
    return h


def approx_hessian(f, x, epsilon):
    """
    Approximate the full Hessian matrix of a function using central
    finite differences

    Parameters
    ----------
    f: callable
      The function to differentiate
    x: ndarray
      Point where the Hessian is to be evaluated
    epsilon: float
      Stepsize for finite differences

    Returns
    -------
    H: ndarray
      Hessian matrix at `x`
    """
    n = len(x)
    H = np.zeros((n, n))
    ei = np.zeros(n)
    for i in range(n):
        ei[i] = .5 * epsilon
        g1 = approx_gradient(f, x + ei, epsilon)
        g2 = approx_gradient(f, x - ei, epsilon)
        H[i, :] = (g1 - g2) / epsilon
        ei[i] = 0
    return H


def smooth(img, sigma):
    """
    Smooth an image by an isotropic Gaussian filter

    Parameters
    ----------
    img: nibabel-like image
      Input image
    sigma: float
      Filter standard deviation in mm

    Returns
    -------
    sdata: ndarray
      Smoothed data array
    """
    if sigma < 0:
        raise ValueError('smoothing kernel size is negative')
    elif sigma == 0:
        return img.get_data()
    else:
        sigma_vox = sigma / np.sqrt(np.sum(img.get_affine()[0:3, 0:3] ** 2, 0))
        return nd.gaussian_filter(img.get_data(), sigma_vox)


def _clamp_array(x, y, bins):

    # Threshold
    dmaxmax = 2 ** (8 * y.dtype.itemsize - 1) - 1
    dmax = bins - 1  # default output maximum value
    if dmax > dmaxmax:
        raise ValueError('Excess number of bins')
    xmin = float(x.min())
    xmax = float(x.max())
    d = xmax - xmin

    """
    If the image dynamic is small, no need for compression: just
    downshift image values and re-estimate the dynamic range (hence
    xmax is translated to xmax-tth casted to the appropriate
    dtype. Otherwise, compress after downshifting image values (values
    equal to the threshold are reset to zero).
    """
    if issubclass(x.dtype.type, np.integer) and d <= dmax:
        y[:] = x - xmin
        bins = int(d) + 1
    else:
        a = dmax / d
        y[:] = np.round(a * (x - xmin))

    return y, bins


def clamp_array(x, bins, mask=None):
    """
    Clamp array values that fall within a given mask in the range
    [0..bins-1] and reset masked values to -1.

    Parameters
    ----------
    x : ndarray
      The input array
    bins : number
      Desired number of bins
    mask : ndarray, tuple or slice
      Anything such that x[mask] is an array.

    Returns
    -------
    y : ndarray
      Clamped array, masked items are assigned -1
    bins : int
      Adjusted number of bins
    """
    if bins > np.iinfo(np.short).max:
        raise ValueError('Too large a bin size')
    y = -np.ones(x.shape, dtype=CLAMP_DTYPE)
    if mask is None:
        y, bins = _clamp_array(x, y, bins)
    else:
        ym = y[mask]
        xm = x[mask]
        ym, bins = _clamp_array(xm, ym, bins)
        y[mask] = ym
    return y, bins


def clamp(img, bins, mask=None, sigma=0):
    """Remap in-mask image intensity values to the range
    [0..bins-1]. Out-of-mask voxels are mapped to -1. A spatial
    Gaussian filter is possibly applied as a pre-processing.

    Parameters
    ----------
    img: nibabel-like image
      Input image
    bins : number
      Desired number of bins
    mask : ndarray, tuple or slice
      Image mask
    sigma: float
      Gaussian filter standard deviation in mm

    Returns
    -------
    data: ndarray
      Clamped data array
    bins: int
      Adjusted number of bins

    """
    data = smooth(img, sigma)
    return clamp_array(data, bins, mask=mask)
