from __future__ import division
import numpy as np
from menpo.base import MenpoMissingDependencyError, partial_doc
from menpo.image.base import normalize_pixels_range
from ..base import winitfeature

try:
    from cyvlfeat.sift.dsift import dsift as cyvlfeat_dsift
except ImportError as e:
    raise MenpoMissingDependencyError(e)


@winitfeature
def dsift(pixels, window_step_horizontal=1, window_step_vertical=1,
          num_bins_horizontal=2, num_bins_vertical=2, num_or_bins=9,
          cell_size_horizontal=6, cell_size_vertical=6, fast=True,
          verbose=False):
    r"""
    Computes a 2-dimensional dense SIFT features image with ``C`` number of
    channels, where
    ``C = num_bins_horizontal * num_bins_vertical * num_or_bins``. The dense
    SIFT [2]_ implementation is taken from Vlfeat [1]_.

    Parameters
    ----------
    pixels : :map:`Image` or subclass or ``(C, Y, X)`` `ndarray`
        Either the image object itself or an array with the pixels. The first
        dimension is interpreted as channels.
    window_step_horizontal : `int`, optional
        Defines the horizontal step by which the window is moved, thus it
        controls the features density. The metric unit is pixels.
    window_step_vertical : `int`, optional
        Defines the vertical step by which the window is moved, thus it
        controls the features density. The metric unit is pixels.
    num_bins_horizontal : `int`, optional
        Defines the number of histogram bins in the X direction.
    num_bins_vertical : `int`, optional
        Defines the number of histogram bins in the Y direction.
    num_or_bins : `int`, optional
        Defines the number of orientation histogram bins.
    cell_size_horizontal : `int`, optional
        Defines cell width in pixels. The cell is the region that is covered by
        a spatial bin.
    cell_size_vertical : `int`, optional
        Defines cell height in pixels. The cell is the region that is covered by
        a spatial bin.
    fast : `bool`, optional
        If ``True``, then the windowing function is a piecewise-flat, rather
        than Gaussian. While this breaks exact SIFT equivalence, in practice it
        is much faster to compute.
    verbose : `bool`, optional
        Flag to print SIFT related information.

    Raises
    ------
    ValueError
        Only 2D arrays are supported
    ValueError
        Size must only contain positive integers.
    ValueError
        Step must only contain positive integers.
    ValueError
        Window size must be a positive integer.
    ValueError
        Geometry must only contain positive integers.

    References
    ----------
    .. [1] Vedaldi, Andrea, and Brian Fulkerson. "VLFeat: An open and portable
       library of computer vision algorithms." Proceedings of the international
       conference on Multimedia. ACM, 2010.
    .. [2] Lowe, David G. "Distinctive image features from scale-invariant
       keypoints." International journal of computer vision 60.2 (2004): 91-110.
    """
    # If norm is set to True, then the centers array will have a third column
    # with descriptor norm, or energy, before contrast normalization.
    # This information can be used to suppress low contrast descriptors.
    centers, output = cyvlfeat_dsift(
        pixels[0], step=[window_step_vertical, window_step_horizontal],
        size=[cell_size_vertical, cell_size_horizontal], bounds=None,
        norm=False, fast=fast, float_descriptors=True,
        geometry=(num_bins_vertical, num_bins_horizontal, num_or_bins),
        verbose=False)

    # the output shape can be calculated from looking at the range of
    # centres / the window step size in each dimension. Note that cyvlfeat
    # returns x, y centres.
    shape = (((centers[-1, :] - centers[0, :]) /
              [window_step_vertical, window_step_horizontal]) + 1)
    shape = shape.astype(np.int)

    # print information
    if verbose:
        info_str = "Dense SIFT features:\n" \
                   "  - Input image is {}W x {}H with {} channels.\n" \
                   "  - Sampling step of ({}W, {}H).\n" \
                   "  - {}W x {}H spatial bins and {} orientation bins.\n" \
                   "  - Cell size of {}W x {}H pixels.\n".format(
                   pixels.shape[2], pixels.shape[1], pixels.shape[0],
                   window_step_horizontal, window_step_vertical,
                   num_bins_horizontal, num_bins_vertical, num_or_bins,
                   cell_size_horizontal, cell_size_vertical)
        if fast:
            info_str += "  - Fast mode is enabled.\n"
        info_str += "Output image size {}W x {}H x {}.".format(
            int(shape[1]), int(shape[0]), output.shape[0])
        print(info_str)

    # return SIFT and centers in the correct form
    return (np.require(np.rollaxis(output.reshape((shape[0], shape[1], -1)),
                                   -1),
                       dtype=np.double, requirements=['C']),
            np.require(centers.reshape((shape[0], shape[1], -1)),
                       dtype=np.int))


# A predefined method for a 'faster' dsift method
fast_dsift = partial_doc(dsift, fast=True, cell_size_vertical=5,
                         cell_size_horizontal=5, num_bins_horizontal=1,
                         num_bins_vertical=1, num_or_bins=8)


# Predefined dsift that returns a 128d vector
def vector_128_dsift(x, dtype=np.float32):
    r"""
    Computes a SIFT feature vector from a square patch (or image). Patch
    **must** be square and the output vector will *always* be a ``(128,)``
    vector. Please see :func:`dsift` for more information.

    Parameters
    ----------
    x : :map:`Image` or subclass or ``(C, Y, Y)`` `ndarray`
        Either the image object itself or an array with the pixels. The first
        dimension is interpreted as channels. Must be square i.e.
        ``height == width``.
    dtype : ``np.dtype``, optional
        The dtype of the returned vector.

    Raises
    ------
    ValueError
        Only square images are supported.
    """
    if not isinstance(x, np.ndarray):
        x = x.pixels
    if x.shape[-1] != x.shape[-2]:
        raise ValueError('This feature only works with square images '
                         'i.e. width == height')
    patch_shape = x.shape[-1]
    n_bins = 4
    c_size = patch_shape // n_bins
    x = normalize_pixels_range(x, error_on_unknown_type=False)
    return dsift(x,
                 window_step_horizontal=patch_shape,
                 window_step_vertical=patch_shape,
                 num_bins_horizontal=n_bins, num_bins_vertical=n_bins,
                 cell_size_horizontal=c_size, cell_size_vertical=c_size,
                 num_or_bins=8, fast=True).astype(dtype)


# Predefined dsift that returns a 128d vector normalized by the hellinger norm
def hellinger_vector_128_dsift(x, dtype=np.float32):
    r"""
    Computes a SIFT feature vector from a square patch (or image). Patch
    **must** be square and the output vector will *always* be a ``(128,)``
    vector. Please see :func:`dsift` for more information.

    The output of :func:`vector_128_dsift` is normalized using the hellinger
    norm (also called the Bhattacharyya distance) which is a measure
    designed to quantify the similarity between two probability distributions.
    Since SIFT is a histogram based feature, this has been shown to improve
    performance. Please see [1]_ for more information.

    Parameters
    ----------
    x : :map:`Image` or subclass or ``(C, Y, Y)`` `ndarray`
        Either the image object itself or an array with the pixels. The first
        dimension is interpreted as channels. Must be square i.e.
        ``height == width``.
    dtype : ``np.dtype``, optional
        The dtype of the returned vector.

    Raises
    ------
    ValueError
        Only square images are supported.

    References
    ----------
    .. [1] Arandjelovic, Relja, and Andrew Zisserman. "Three things everyone
           should know to improve object retrieval.", CVPR, 2012.
    """
    h = vector_128_dsift(x, dtype=dtype)
    h /= (h.sum(axis=0) + 1e-15)
    return np.sqrt(h)
