"""
..
    Copyright (c) 2014, Magni developers.
    All rights reserved.
    See LICENSE.rst for further information.

Module providing AFM image reconstruction and analysis of reconstructed images.

Routine listings
----------------
analyse(x, Phi, Psi)
    Sample an image, reconstruct it, and analyse the reconstructed image.
reconstruct(y, Phi, Psi)
    Reconstruct an image from compressively sensed measurements.

"""

from __future__ import division

import numpy as np

from magni.afm import config as _config
from magni.cs.reconstruction import iht as _iht
from magni.cs.reconstruction import sl0 as _sl0
from magni.imaging import evaluation as _eval
from magni.utils.matrices import Matrix as _Matrix
from magni.utils.matrices import MatrixCollection as _MatrixC
from magni.utils.validation import decorate_validation as _decorate_validation
from magni.utils.validation import validate_ndarray as _validate_ndarray


@_decorate_validation
def _validate_analyse(x, Phi, Psi):
    """
    Validate the `analyse` function.

    See Also
    --------
    analyse : The validated function.
    magni.utils.validation.validate_ndarray : Validation.
    magni.utils.validation.validate_matrix : Validation.

    """

    try:
        _validate_ndarray(Phi, 'Phi')
    except TypeError:
        if not isinstance(Phi, _Matrix) and not isinstance(Phi, _MatrixC):
            raise TypeError('Phi must be a matrix.')

    try:
        _validate_ndarray(Psi, 'Psi')
    except TypeError:
        if not isinstance(Psi, _Matrix) and not isinstance(Psi, _MatrixC):
            raise TypeError('Psi must be a matrix.')

    if Phi.shape[1] != Psi.shape[0]:
        raise ValueError('Phi and Psi must have compatible shapes.')

    _validate_ndarray(x, 'x', {'shape': (Phi.shape[1], 1)})


def analyse(x, Phi, Psi):
    """
    Sample an image, reconstruct it, and analyse the reconstructed image.

    Parameters
    ----------
    x : numpy.ndarray
        The original image vector.
    Phi : magni.utils.matrices.Matrix or numpy.matrix
        The measurement matrix.
    Psi : magni.utils.matrices.Matrix or numpy.matrix
        The dictionary.

    Returns
    -------
    x : numpy.ndarray
        The reconstructed image vector.

    See Also
    --------
    magni.afm.config : Configuration options.
    magni.imaging.evaluation : Image reconstruction quality evaluation.

    Examples
    --------
    Prior to the actual example, data is loaded and a measurement matrix and a
    dictionary are defined. First, the example MI file provided with the
    package is loaded:

    >>> from magni.afm.reconstruction import analyse
    >>> path = magni.utils.split_path(magni.__path__[0])[0]
    >>> path = path + 'examples' + os.sep + 'example.mi'
    >>> mi_file = magni.afm.io.read_mi_file(path)
    >>> mi_buffer = mi_file.get_buffer('Topography')[0]
    >>> mi_data = mi_buffer.get_data()
    >>> x = magni.imaging.mat2vec(mi_data)

    Next, a measurement matrix is defined. This matrix is equal to the matrix
    generated by running `np.eye(len(x))[::2, :]` but for speed, the matrix is
    instead defined with fast operations:

    >>> def Phi_A(x):
    ...     y = x[::2]
    ...     return y
    >>> def Phi_T(y):
    ...     x = np.zeros((2 * len(y), 1))
    ...     x[::2] = y
    ...     return x
    >>> Phi = magni.utils.matrices.Matrix(Phi_A, Phi_T, (),
    ...                                   (int(len(x) / 2), len(x)))

    Next, a dictionary is defined. This dictionary is the DCT basis likewise
    defined with fast operations:

    >>> Psi = magni.imaging.dictionaries.get_DCT(mi_data.shape)

    Finally, the actual example:

    >>> print('MSE: {:.2f}, PSNR: {:.2f}'.format(*analyse(x, Phi, Psi)))
    MSE: 0.23, PSNR: 5.92

    """

    _validate_analyse(x, Phi, Psi)

    y = Phi.dot(x)
    x_hat = reconstruct(y, Phi, Psi)

    mse = _eval.calculate_mse(x, x_hat)
    psnr = _eval.calculate_psnr(x, x_hat, float(np.abs(x).max()))

    return (mse, psnr)


@_decorate_validation
def _validate_reconstruct(y, Phi, Psi):
    """
    Validate the `reconstruct` function.

    See Also
    --------
    reconstruct : The validated function.
    magni.utils.validation.validate_ndarray : Validation.
    magni.utils.validation.validate_matrix : Validation.

    """

    try:
        _validate_ndarray(Phi, 'Phi')
    except TypeError:
        if not isinstance(Phi, _Matrix) and not isinstance(Phi, _MatrixC):
            raise TypeError('Phi must be a matrix.')

    try:
        _validate_ndarray(Psi, 'Psi')
    except TypeError:
        if not isinstance(Psi, _Matrix) and not isinstance(Psi, _MatrixC):
            raise TypeError('Psi must be a matrix.')

    if Phi.shape[1] != Psi.shape[0]:
        raise ValueError('Phi and Psi must have compatible shapes.')

    _validate_ndarray(y, 'y', {'shape': (Phi.shape[0], 1)})


def reconstruct(y, Phi, Psi):
    """
    Reconstruct an image from compressively sensed measurements.

    Parameters
    ----------
    y : numpy.ndarray
        The measurement vector.
    Phi : magni.utils.matrices.Matrix or numpy.matrix
        The measurement matrix.
    Psi : magni.utils.matrices.Matrix or numpy.matrix
        The dictionary.

    Returns
    -------
    x : numpy.ndarray
        The reconstructed image vector.

    See Also
    --------
    magni.afm.config : Configuration options.
    magni.cs.reconstruction : Compressed sensing reconstruction algorithms.

    Examples
    --------
    Prior to the actual example, data is loaded and a measurement matrix and a
    dictionary are defined. First, the example MI file provided with the
    package is loaded:

    >>> from magni.afm.reconstruction import reconstruct
    >>> path = magni.utils.split_path(magni.__path__[0])[0]
    >>> path = path + 'examples' + os.sep + 'example.mi'
    >>> mi_file = magni.afm.io.read_mi_file(path)
    >>> mi_buffer = mi_file.get_buffer('Topography')[0]
    >>> mi_data = mi_buffer.get_data()
    >>> x = magni.imaging.mat2vec(mi_data)

    Next, a measurement matrix is defined. This matrix is equal to the matrix
    generated by running `np.eye(len(x))[::2, :]` but for speed, the matrix is
    instead defined with fast operations:

    >>> def Phi_A(x):
    ...     y = x[::2]
    ...     return y
    >>> def Phi_T(y):
    ...     x = np.zeros((2 * len(y), 1))
    ...     x[::2] = y
    ...     return x
    >>> Phi = magni.utils.matrices.Matrix(Phi_A, Phi_T, (),
    ...                                   (int(len(x) / 2), len(x)))

    Next, a dictionary is defined. This dictionary is the DCT basis likewise
    defined with fast operations:

    >>> Psi = magni.imaging.dictionaries.get_DCT(mi_data.shape)

    Finally, the actual example:

    >>> y = Phi.dot(x)
    >>> print('Maximum absolute pixel error: {:.3f}'
    ...       .format(np.abs(reconstruct(y, Phi, Psi) - x).max()))
    Maximum absolute pixel error: 0.960

    """

    _validate_reconstruct(y, Phi, Psi)

    A = _MatrixC((Phi, Psi))
    algorithm = _config.get('algorithm')

    if algorithm == 'iht':
        algorithm = _iht.run
    elif algorithm == 'sl0':
        A = A.A
        algorithm = _sl0.run

    alpha = algorithm(y, A)
    x = Psi.dot(alpha)
    return x
