"""
ffit.py
Functions for Fitness

Lev 1 Feb 2009
"""

import scipy.optimize
import numpy

class CONST(numpy.float64):
    """
    This class is intended be used to fix a parameter during a fit.
    """
    pass

def leastsq(x, y, func, params, mask = None, yerr = None, Var = None, **vargs):
    """
    A frontend to scipy least square fit.
    
    'x' and 'y' are the data feeded to a fit. If 'y' is None, 'func(x, *params)'
    is used as residuals. Otherwise 'sum((y - func(x, *params))**2)' is minimized, e.g.
    'func(x, *params)' is a model for data rather than residuals.
    See 'yerr' and 'Var' below for more details.
    
    Only a subset of 'params' is varied as specified by the 'mask'.
    
    'params' is a mixture of initial guesses of the variated parameters and
    values of the fixed ones.

    'mask' is a boolean array of the same size as 'params'. If 'mask[n]' is True,
    the 'params[n]' is variated, otherwise 'params[n]' is fixed. If 'mask' is
    not specified it is deduced from the 'params' itself. A parameter
    is considered fixed if it is of type 'ffit.CONST' derived from numpy.float64.
    
    'yerr' is a numerical array that can be used to weight the samples.
    If both 'y' and 'yerr' are not 'None', 'sum((y - func(x, *params))**2 / yerr**2)' is minimised.
    
    'Var' is a 2 dimensional numerical array that can be used to introduce
    co-variance to the samples. If 'Var' is supplied, 'yerr' is ignored and
    the 'sum(abs(y - func(x, *params)))'
    
    Any extra keyword arguments are forwarded to the 'scipy.optimize.leastsq'.

    The return value has the same format as for 'scipy.optimize.leastsq',
    except for the array of fit results (first element of the returned
    tuple): here it contains all 'params', both varied and fixed in the same
    order as in the input argument. The rest of the tuple depends on
    'full_output' keyword argument.
    
    Example:
    
    >>> def func(x, A, B): return A*x + B
    ...
    >>> x = numpy.array([1,2,3])
    >>> y = numpy.array([1,2,3])
    >>> res = ffit.leastsq(x, y, func, [1, 1], [True, True])
    >>> print res[0]
    [  1.00000000e+00  -4.06333240e-09]
    >>> res = ffit.leastsq(x, y, func, [1, 0.25], [True, False])
    >>> print res[0]
    [ 0.89285714  0.25      ]
    >>> res = ffit.leastsq(x, y, func, [ffit.CONST(0.25), 1])
    >>> print res[0]
    [ 0.25  1.5 ]

    In the second example the intercept of the linear fit is kept fixed at 1,
    the best fit being y = 0.89*x + 1, rather than y = 1*x + 0.25.
    Similarily in the third one the slope is fixed at 0.25, and the best ift
    is y = 0.25*x + 1.5.
    """
    # if 'mask' was not specified, create one with 'True' for every element of
    # 'params' array except for those of 'fit.const' type.
    if mask is None:
        mask = [not isinstance(p, CONST) for p in params]
#        print mask
    mask = numpy.asarray(mask, dtype=bool)    
    
    # duplicate params, because the fit results are filled into it after the fit
    params = numpy.array(params, dtype=numpy.float64, copy=True)
    if params.ndim != 1:
        raise ValueError("'params' must be a 1D array")
    
    if mask.shape != params.shape:
        raise ValueError("'params' and 'mask' must be the same size")

    if not numpy.any(mask):
        raise ValueError("all 'params' are constrained. no degrees of freedom for a fit")

    if y is None:
        # build a residual calculation function
        def residuals(fitparams, mask, x, consts):
            # combine 'fitparams' and 'consts' into 'params'
            params = numpy.zeros(len(mask))
            params[mask] = fitparams
            params[~mask] = consts
            return func(x, *params)
        # fit
        res = scipy.optimize.leastsq(residuals, params[mask], (mask, x, params[~mask]), **vargs)
    else:
        # build a residual calculation function
        x = numpy.asarray(x)
        y = numpy.asarray(y)
        if y.shape != x.shape:
            raise ValueError("'x' and 'y' are different size")
        if Var is None:
            if yerr is None:
                yerr = numpy.ones(y.shape)
            else:
                yerr = numpy.asarray(yerr)
                if y.shape != yerr.shape:
                    raise ValueError("'y' and 'yerr' are different size")
            yerr_arg = yerr
            def residuals(fitparams, mask, x, y, yerr, consts):
                # combine 'fitparams' and 'consts' into 'params'
                params = numpy.zeros(len(mask))
                params[mask] = fitparams
                params[~mask] = consts
                return (y - func(x, *params)) / yerr
        else:
            Var = numpy.asarray(Var)
            if len(Var.shape) != 2 or Var.shape[0] != len(x) or Var.shape[1] != len(x):
                raise ValueError("'Var' should be square 'len(x)' by 'len(x)' in size")
            invVar = numpy.matrix(Var).I
            yerr_arg = invVar
            def residuals(fitparams, mask, x, y, invVar, consts):
                # combine 'fitparams' and 'consts' into 'params'
                params = numpy.zeros(len(mask))
                params[mask] = fitparams
                params[~mask] = consts
                res = numpy.matrix(y - func(x, *params))
                return numpy.ones(x.shape) * (res * invVar * res.T)[0,0]**0.5
        # fit
        res = scipy.optimize.leastsq(residuals, params[mask], (mask, x, y, yerr_arg, params[~mask]), **vargs)

    # replace initial guesses with the fit results, fixed parameters are already in place.
    params[mask] = res[0]
    res = tuple([params] + list(res[1:]))
    return res