"""
Python wrappers for Orthogonal Distance Regression (ODRPACK).

Notes
=====

* Array formats -- FORTRAN stores its arrays in memory column first, i.e. an
  array element A(i, j, k) will be next to A(i+1, j, k). In C and, consequently,
  NumPy, arrays are stored row first: A[i, j, k] is next to A[i, j, k+1]. For
  efficiency and convenience, the input and output arrays of the fitting
  function (and its Jacobians) are passed to FORTRAN without transposition.
  Therefore, where the ODRPACK documentation says that the X array is of shape
  (N, M), it will be passed to the Python function as an array of shape (M, N).
  If M==1, the one-dimensional case, then nothing matters; if M>1, then your
  Python functions will be dealing with arrays that are indexed in reverse of
  the ODRPACK documentation. No real biggie, but watch out for your indexing of
  the Jacobians: the i,j'th elements (@f_i/@x_j) evaluated at the n'th
  observation will be returned as jacd[j, i, n]. Except for the Jacobians, it
  really is easier to deal with x[0] and x[1] than x[:,0] and x[:,1]. Of course,
  you can always use the transpose() function from scipy explicitly.

* Examples -- See the accompanying file test/test.py for examples of how to set
  up fits of your own. Some are taken from the User's Guide; some are from
  other sources.

* Models -- Some common models are instantiated in the accompanying module
  models.py . Contributions are welcome.

Credits
=======

* Thanks to Arnold Moene and Gerard Vermeulen for fixing some killer bugs.

Robert Kern
robert.kern@gmail.com

"""

from __future__ import division, print_function, absolute_import

import numpy
from warnings import warn
from scipy.odr import __odrpack

__all__ = ['odr', 'OdrWarning', 'OdrError', 'OdrStop',
           'Data', 'RealData', 'Model', 'Output', 'ODR',
           'odr_error', 'odr_stop']

odr = __odrpack.odr


class OdrWarning(UserWarning):
    """
    Warning indicating that the data passed into
    ODR will cause problems when passed into 'odr'
    that the user should be aware of.
    """
    pass


class OdrError(Exception):
    """
    Exception indicating an error in fitting.

    This is raised by `~scipy.odr.odr` if an error occurs during fitting.
    """
    pass


class OdrStop(Exception):
    """
    Exception stopping fitting.

    You can raise this exception in your objective function to tell
    `~scipy.odr.odr` to stop fitting.
    """
    pass


# Backwards compatibility
odr_error = OdrError
odr_stop = OdrStop

__odrpack._set_exceptions(OdrError, OdrStop)


def _conv(obj, dtype=None):
    """ Convert an object to the preferred form for input to the odr routine.
    """

    if obj is None:
        return obj
    else:
        if dtype is None:
            obj = numpy.asarray(obj)
        else:
            obj = numpy.asarray(obj, dtype)
        if obj.shape == ():
            # Scalar.
            return obj.dtype.type(obj)
        else:
            return obj


def _report_error(info):
    """ Interprets the return code of the odr routine.

    Parameters
    ----------
    info : int
        The return code of the odr routine.

    Returns
    -------
    problems : list(str)
        A list of messages about why the odr() routine stopped.
    """

    stopreason = ('Blank',
                  'Sum of squares convergence',
                  'Parameter convergence',
                  'Both sum of squares and parameter convergence',
                  'Iteration limit reached')[info % 5]

    if info >= 5:
        # questionable results or fatal error

        I = (info//10000 % 10,
             info//1000 % 10,
             info//100 % 10,
             info//10 % 10,
             info % 10)
        problems = []

        if I[0] == 0:
            if I[1] != 0:
                problems.append('Derivatives possibly not correct')
            if I[2] != 0:
                problems.append('Error occurred in callback')
            if I[3] != 0:
                problems.append('Problem is not full rank at solution')
            problems.append(stopreason)
        elif I[0] == 1:
            if I[1] != 0:
                problems.append('N < 1')
            if I[2] != 0:
                problems.append('M < 1')
            if I[3] != 0:
                problems.append('NP < 1 or NP > N')
            if I[4] != 0:
                problems.append('NQ < 1')
        elif I[0] == 2:
            if I[1] != 0:
                problems.append('LDY and/or LDX incorrect')
            if I[2] != 0:
                problems.append('LDWE, LD2WE, LDWD, and/or LD2WD incorrect')
            if I[3] != 0:
                problems.append('LDIFX, LDSTPD, and/or LDSCLD incorrect')
            if I[4] != 0:
                problems.append('LWORK and/or LIWORK too small')
        elif I[0] == 3:
            if I[1] != 0:
                problems.append('STPB and/or STPD incorrect')
            if I[2] != 0:
                problems.append('SCLB and/or SCLD incorrect')
            if I[3] != 0:
                problems.append('WE incorrect')
            if I[4] != 0:
                problems.append('WD incorrect')
        elif I[0] == 4:
            problems.append('Error in derivatives')
        elif I[0] == 5:
            problems.append('Error occurred in callback')
        elif I[0] == 6:
            problems.append('Numerical error detected')

        return problems

    else:
        return [stopreason]


class Data(object):
    """
    The data to fit.

    Parameters
    ----------
    x : array_like
        Observed data for the independent variable of the regression
    y : array_like, optional
        If array-like, observed data for the dependent variable of the
        regression. A scalar input implies that the model to be used on
        the data is implicit.
    we : array_like, optional
        If `we` is a scalar, then that value is used for all data points (and
        all dimensions of the response variable).
        If `we` is a rank-1 array of length q (the dimensionality of the
        response variable), then this vector is the diagonal of the covariant
        weighting matrix for all data points.
        If `we` is a rank-1 array of length n (the number of data points), then
        the i'th element is the weight for the i'th response variable
        observation (single-dimensional only).
        If `we` is a rank-2 array of shape (q, q), then this is the full
        covariant weighting matrix broadcast to each observation.
        If `we` is a rank-2 array of shape (q, n), then `we[:,i]` is the
        diagonal of the covariant weighting matrix for the i'th observation.
        If `we` is a rank-3 array of shape (q, q, n), then `we[:,:,i]` is the
        full specification of the covariant weighting matrix for each
        observation.
        If the fit is implicit, then only a positive scalar value is used.
    wd : array_like, optional
        If `wd` is a scalar, then that value is used for all data points
        (and all dimensions of the input variable). If `wd` = 0, then the
        covariant weighting matrix for each observation is set to the identity
        matrix (so each dimension of each observation has the same weight).
        If `wd` is a rank-1 array of length m (the dimensionality of the input
        variable), then this vector is the diagonal of the covariant weighting
        matrix for all data points.
        If `wd` is a rank-1 array of length n (the number of data points), then
        the i'th element is the weight for the i'th input variable observation
        (single-dimensional only).
        If `wd` is a rank-2 array of shape (m, m), then this is the full
        covariant weighting matrix broadcast to each observation.
        If `wd` is a rank-2 array of shape (m, n), then `wd[:,i]` is the
        diagonal of the covariant weighting matrix for the i'th observation.
        If `wd` is a rank-3 array of shape (m, m, n), then `wd[:,:,i]` is the
        full specification of the covariant weighting matrix for each
        observation.
    fix : array_like of ints, optional
        The `fix` argument is the same as ifixx in the class ODR. It is an
        array of integers with the same shape as data.x that determines which
        input observations are treated as fixed. One can use a sequence of
        length m (the dimensionality of the input observations) to fix some
        dimensions for all observations. A value of 0 fixes the observation,
        a value > 0 makes it free.
    meta : dict, optional
        Free-form dictionary for metadata.

    Notes
    -----
    Each argument is attached to the member of the instance of the same name.
    The structures of `x` and `y` are described in the Model class docstring.
    If `y` is an integer, then the Data instance can only be used to fit with
    implicit models where the dimensionality of the response is equal to the
    specified value of `y`.

    The `we` argument weights the effect a deviation in the response variable
    has on the fit.  The `wd` argument weights the effect a deviation in the
    input variable has on the fit. To handle multidimensional inputs and
    responses easily, the structure of these arguments has the n'th
    dimensional axis first. These arguments heavily use the structured
    arguments feature of ODRPACK to conveniently and flexibly support all
    options. See the ODRPACK User's Guide for a full explanation of how these
    weights are used in the algorithm. Basically, a higher value of the weight
    for a particular data point makes a deviation at that point more
    detrimental to the fit.

    """

    def __init__(self, x, y=None, we=None, wd=None, fix=None, meta={}):
        self.x = _conv(x)

        if not isinstance(self.x, numpy.ndarray):
            raise ValueError(("Expected an 'ndarray' of data for 'x', "
                              "but instead got data of type '{name}'").format(
                    name=type(self.x).__name__))

        self.y = _conv(y)
        self.we = _conv(we)
        self.wd = _conv(wd)
        self.fix = _conv(fix)
        self.meta = meta

    def set_meta(self, **kwds):
        """ Update the metadata dictionary with the keywords and data provided
        by keywords.

        Examples
        --------
        ::

            data.set_meta(lab="Ph 7; Lab 26", title="Ag110 + Ag108 Decay")
        """

        self.meta.update(kwds)

    def __getattr__(self, attr):
        """ Dispatch attribute access to the metadata dictionary.
        """
        if attr in self.meta:
            return self.meta[attr]
        else:
            raise AttributeError("'%s' not in metadata" % attr)


class RealData(Data):
    """
    The data, with weightings as actual standard deviations and/or
    covariances.

    Parameters
    ----------
    x : array_like
        Observed data for the independent variable of the regression
    y : array_like, optional
        If array-like, observed data for the dependent variable of the
        regression. A scalar input implies that the model to be used on
        the data is implicit.
    sx : array_like, optional
        Standard deviations of `x`.
        `sx` are standard deviations of `x` and are converted to weights by
        dividing 1.0 by their squares.
    sy : array_like, optional
        Standard deviations of `y`.
        `sy` are standard deviations of `y` and are converted to weights by
        dividing 1.0 by their squares.
    covx : array_like, optional
        Covariance of `x`
        `covx` is an array of covariance matrices of `x` and are converted to
        weights by performing a matrix inversion on each observation's
        covariance matrix.
    covy : array_like, optional
        Covariance of `y`
        `covy` is an array of covariance matrices and are converted to
        weights by performing a matrix inversion on each observation's
        covariance matrix.
    fix : array_like, optional
        The argument and member fix is the same as Data.fix and ODR.ifixx:
        It is an array of integers with the same shape as `x` that
        determines which input observations are treated as fixed. One can
        use a sequence of length m (the dimensionality of the input
        observations) to fix some dimensions for all observations. A value
        of 0 fixes the observation, a value > 0 makes it free.
    meta : dict, optional
        Free-form dictionary for metadata.

    Notes
    -----
    The weights `wd` and `we` are computed from provided values as follows:

    `sx` and `sy` are converted to weights by dividing 1.0 by their squares.
    For example, ``wd = 1./numpy.power(`sx`, 2)``.

    `covx` and `covy` are arrays of covariance matrices and are converted to