"""
Least-square fits of non-linear models with up to O(10^3 - 10^4) pars.

This module contains utility classes to perform least square fits of
(potentially large) non-linear models to experimental data.  The two
main features implemented are

- `Chi2`: a class that take a model and a dataset as arguments and
  computes the value, gradient and Hessian of the chi2 function. The
  computation tries to be efficient by using sparse matrices. This
  yield very significant speed-ups if the problem is sparse.

- `fmin_ncg`: a rewrite of scipy.optimize.fmin_ncg that accepts sparse
Jacobian and Hessian and take specific care to the factorize the
computation of the chi2 value and its derivatives.

A chi2 object may still be passed to any minimizer, notably those in
``scipy.optimize''.
"""

import logging

import numpy as np
import scipy.optimize
import scipy.sparse.linalg
import selinv
from sksparse.cholmod import cholesky, cholesky_AAt

_epsilon = 1.4901161193847656e-08


class Chi2(object):
    def __init__(self, model, x, y, w=None, s=None):
        r"""Chi2 funtion of a non-linear model of the form:

        $$\chi^2 = R^t R$$
        with:
        $$R = w (y - model(x, p))$$

        Attention: w is the square root of the inverse of the covariance
        matrix of y. At the moment we only support diagonal covariance
        matrices, specified by their diagonal as a 1d-array because I
        made specific optimizations for this case.

        Parameters:
        -----------
        model: callable
               call signature must be model(x, p, jac) and the return value must be either
               - model(x, p) if jac=False
               - (model(x,p), dmodel / dp) if jac is True, and dmodel/dp must be a sparse coo matrix
        x: x is passed directly as an argument to model
        y: 1d array
           The measurement. If si or w are given the data will be weighted
        s: array or None (broadcastable to the shape of y) expected to be the 1-sigma uncertainty of the measurements
        w: array or None (broadcastable to the shape of y) expected to be the inverse of s.
        """
        self.model = model
        self.xi = x
        self.yi = y
        assert w is None or s is None, "Cannot provide w and s at the same time"

        self.w = None
        if w is not None:
            self.w = w
        if s is not None:
            self.w = 1.0 / s

    def __call__(self, p, grad=False, hess=False, goods=None):
        # Compute the model and the Jacobian if needed
        if grad or hess:
            m, J = self.model(self.xi, p, jac=True)
            if self.w is not None:
                J.data *= self.w[J.row]
            J = J.tocsr()
            if goods is not None:
                J = J[goods, :]
        else:
            m = self.model(self.xi, p, jac=False)
        # Weight the residuals if weights are available
        wres = self.yi - m
        if self.w is not None:
            wres *= self.w
        if goods is not None:
            wres = wres[goods]

        # Chi2
        ret = [wres.T.dot(wres)]
        # print ret
        if not (grad or hess):
            return ret[0]
        if grad:
            ret.append(-2.0 * J.T.dot(wres))
        if hess:
            ret.append(2.0 * J.T * J)

        return ret

    def get_res(self, p, goods=None):
        m = self.model(self.xi, p, jac=False)
        wres = self.yi - m
        if self.w is not None:
            wres *= self.w
        if goods is not None:
            return wres[goods]
        else:
            return wres

    def gradient(self, p, **keys):
        return self(p, grad=True, **keys)[1]

    def hessian(self, p, **keys):
        return self(p, hess=True, **keys)[1]

    def densehessian(self, p, **keys):
        return self(p, hess=True, **keys)[1].todense()

    def get_cov(self, p, **keys):
        H = self.hessian(p, **keys)
        C = scipy.sparse.linalg.inv(H)
        return 2 * C

    def get_diag_cov(self, p, **keys):
        H = self.hessian(p, **keys)
        if hasattr(self, "factor"):
            return 2 * selinv.selinv(H, self.factor.P())
        else:
            return 2 * selinv.selinv(H)

    def minimize(
        self,
        p0,
        algorithm="fmin_ncg",
        goods=None,
        maxiter=2e3,
        fallback=True,
        verbose=True,
        **keys,
    ):
        if algorithm == "fmin_ncg":
            return fmin_ncg(self, p0, self.gradient, fhess=self.hessian, **keys)
        elif algorithm == "simplex":
            return scipy.optimize.fmin(self, p0, **keys)
        elif algorithm == "slsqp":
            return scipy.optimize.fmin_slsqp(self, p0, fprime=self.gradient, **keys)
        elif algorithm == "gauss_newton_homebrew":
            # Gauss-Newton with a confidence region check to switch to
            # steepest decent in case of trouble. Tipycally fails to
            # converge in high-residual cases. Try Newton or
            # quasi-Newton in this cases.
            raise_exception = keys.pop("raise", False)
            x = p0.copy()
            epsilon4 = 0.5
            it = 0
            # if "maxiter" in keys.keys():
            #     maxiter = keys["maxiter"]
            # else:
            #     maxiter = 2e3
            # if "fallback" in keys.keys():
            #     fallback = keys["fallback"]
            # else:
            #     fallback = True

            fup = 2.0
            while True:
                m, J = self.model(self.xi, x, jac=True)
                if self.w is not None:
                    J.data *= self.w[J.row]
                J = J.tocsr()
                if goods is not None:
                    J = J[goods, :]
                    m = m[goods]
                    yi = self.yi[goods]
                else:
                    yi = self.yi
                wres = yi - m
                if self.w is not None:
                    if goods is not None:
                        wres *= self.w[goods]
                    else:
                        wres *= self.w
                chi2 = wres.T.dot(wres)  # Why not call Chi2.__call__ ?
                g = J.T.dot(wres)  # SZF why not -2 * this ?
                if it == 0:
                    self.factor = cholesky_AAt(J.T)
                else:
                    self.factor.cholesky_AAt_inplace(J.T)
                deltax = self.factor(g).squeeze()
                chi2new = self(x + deltax, grad=False, hess=False, goods=goods)
                actual_red = chi2 - chi2new
                # print np.linalg.cond((J.T * J).todense())
                # For a linear f:
                pred_red = np.dot(deltax, g)
                rho = actual_red / pred_red
                dof = len(yi) - len(x)
                if verbose:
                    print(
                        (
                            "chi2=%.2f (dof=%.2f) -> %.2f (%.2f / %.2f = %g), |g|=%g"
                            % (
                                chi2,
                                dof,
                                chi2new,
                                actual_red,
                                pred_red,
                                rho,
                                np.sqrt(np.dot(g, g)),
                            )
                        )
                    )
                if (
                    not (fallback) or rho > 0
                ):  # accept the Gauss-Newton step even if it is rather poor
                    x += deltax
                    # Convergence criterion
                    if actual_red < 0.01:  # and pred_red < 0.01 and rho > 0.9:
                        break
                elif fallback and rho <= 0:  # attempt to save the meuble
                    if verbose:
                        print("bad step: trying steepest decent")

                    def c(mu):
                        return self(x + mu * g, grad=False, hess=False, goods=goods)

                    mu, chi2new, niter, funcalls = scipy.optimize.brent(
                        c, brack=(0, 1), tol=1e-2, full_output=True
                    )
                    if verbose:
                        print(
                            (
                                "Brent found new minimum in %d iter: x + %e * g -> %.2f"
                                % (niter, mu, chi2new)
                            )
                        )
                    if mu < 0:
                        if raise_exception:
                            mu = np.linspace(-1e-4, 1e-4, 20)
                            c = [
                                self(x + _mu * g, grad=False, hess=False) for _mu in mu
                            ]
                            import matplotlib.pyplot as plt

                            plt.plot(mu, c)
                            plt.show()
                            raise ValueError("Dear Lord, help us out of this merdier")
                        else:
                            pass
                    x += mu * g

                    # print "bad step"
                    # d = np.bincount(J.indices, J.data * J.data)
                    # #deltax = g / d
                    # H = J.T * J
                    # d = scipy.sparse.diags(d, 0).tocsc()
                    # nu = 0.1
                    # while (chi2new > chi2):
                    #     self.factor.cholesky_inplace(J.T*J + nu * d)
                    #     #self.factor.cholesky_AAt_inplace(J.T, nu)
                    #     deltax = self.factor(g).squeeze()
                    #     chi2new = self(x + deltax, grad=False, hess=False)
                    #     print chi2new
                    #     nu *= 2
                    #     if nu > 1e2:
                    #         # Holly shit we are in a bad shape
                    #         mu = np.linspace(0,1e-4,10)
                    #         c = [self(x + _mu * g, grad=False, hess=False) for _mu in mu]
                    #         import matplotlib.pyplot as plt
                    #         plt.plot(mu, c)
                    #         plt.show()
                    #         raise ValueError("Dear Lord, help us out of this merdier")
                    #         #x2 = scipy.optimize.anneal(self, x)
                    #         #deltax = x2 - x
                    # print "smoothing param: %g (chi2: %.2f -> %.2f)" % (nu, chi2, chi2new)
                    # x += deltax
                    # #x = scipy.optimize.fmin(self, x)

                it = it + 1
                if it > maxiter:
                    raise ValueError("Grosse fatigue")
                    if verbose:
                        print("Reach maximum number of iteration")
                    break
            chi2 = self(x, goods=goods)
            dof = len(yi) - len(x)
            if verbose:
                print(
                    (
                        "Converged in %d iterations: chi2=%g, d.o.f.=%d, chi2/dof=%.3f"
                        % (it, chi2, dof, chi2 / dof)
                    )
                )
            if keys.pop("full_output", False):
                return x, chi2, dof, self.factor
            else:
                return x

        elif algorithm == "newton_homebrew":
            # Gauss-Newton with a switch to full newton when steps fail. Requires the model to be able to compute
            # the second order part the hessian given the
            # residuals.
            raise_exception = keys.pop("raise", False)
            x = p0.copy()
            it = 0
            while True:
                m, J = self.model(self.xi, x, jac=True)
                if self.w is not None:
                    J.data *= self.w[J.row]
                J = J.tocsr()
                if goods is not None:
                    J = J[goods, :]
                    m = m[goods]
                    yi = self.yi[goods]
                else:
                    yi = self.yi
                wres = yi - m
                if self.w is not None:
                    if goods is not None:
                        wres *= self.w[goods]
                    else:
                        wres *= self.w
                chi2 = wres.T.dot(wres)
                g = J.T.dot(wres)
                H = J.T * J
                if it == 0:
                    self.factor = cholesky(H)
                else:
                    self.factor.cholesky_inplace(H)
                deltax = self.factor(g).squeeze()
                chi2new = self(x + deltax, grad=False, hess=False, goods=goods)
                actual_red = chi2 - chi2new
                # print np.linalg.cond((J.T * J).todense())
                # For a linear f:
                pred_red = np.dot(deltax, g)
                rho = actual_red / pred_red
                dof = len(yi) - len(x)
                print(
                    (
                        "%.2f (%.2f) -> %.2f (%.2f / %.2f = %g), |g|=%g"
                        % (
                            chi2,
                            dof,
                            chi2new,
                            actual_red,
                            pred_red,
                            rho,
                            np.sqrt(np.dot(g, g)),
                        )
                    )
                )
                # accept the Gauss-Newton step even if it is rather poor
                if rho >= 0:
                    x += deltax
                # Convergence criterion
                if actual_red < 0.01 and pred_red < 0.01:  # and rho >= 0:
                    break
                if 0:  # attempt a full newton step
                    Q = self.model.hess(x, wres)
                    H = H + Q
                    self.factor = cholesky(H)
                    deltax = self.factor(g).squeeze()
                    R = g - H * deltax
                    delta2x = self.factor(R)
                    print((abs(delta2x / deltax).max()))
                    chi2new = self(x + deltax, grad=False, hess=False)
                    print(chi2new)
                    chi2new = self(x + deltax + delta2x, grad=False, hess=False)
                    print(chi2new)
                    actual_red = chi2 - chi2new
                    pred_red = np.dot(deltax, g)
                    rho = actual_red / pred_red
                    print(
                        (
                            "%.2f (%.2f) -> %.2f (%.2f / %.2f = %g), |g|=%g"
                            % (
                                chi2,
                                dof,
                                chi2new,
                                actual_red,
                                pred_red,
                                rho,
                                np.sqrt(np.dot(g, g)),
                            )
                        )
                    )
                    if rho < 0:
                        print("Newton step failed. Stopping")
                        break
                    else:
                        x += deltax
                        if actual_red < 0.01 and pred_red < 0.01 and rho > 0.9:
                            print("optimization terminated succesfully")
                            break

                it = it + 1
                if it > 2e3:
                    raise ValueError("Grosse fatigue")
            chi2 = self(x, goods=goods)
            dof = len(yi) - len(x)
            print(
                (
                    "Converged in %d iterations: chi2=%g, d.o.f.=%d, chi2/dof=%.3f"
                    % (it, chi2, dof, chi2 / dof)
                )
            )
            if keys.pop("full_output", False):
                return x, chi2, dof, self.factor
            else:
                return x

        elif algorithm == "gauss_newton_dense":
            # Gauss-Newton with a confidence region check to switch to
            # steepest decent in case of trouble. Tipycally fails to
            # converge in high-residual cases. Try Newton or
            # quasi-Newton in this cases.
            raise_exception = keys.pop("raise", False)
            x = p0.copy()
            # epsilon4 = 0.5
            it = 0
            # if "maxiter" in keys.keys():
            #     maxiter = keys["maxiter"]
            # else:
            #     maxiter = 2e3
            # if "fallback" in keys.keys():
            #     fallback = keys["fallback"]
            # else:
            #     fallback = True

            # fup = 2.0
            while True:
                m, J = self.model(self.xi, x, jac=True)
                if self.w is not None:
                    J.data *= self.w[J.row]
                J = J.tocsr()
                if goods is not None:
                    J = J[goods, :]
                    m = m[goods]
                    yi = self.yi[goods]
                else:
                    yi = self.yi
                wres = yi - m
                if self.w is not None:
                    if goods is not None:
                        wres *= self.w[goods]
                    else:
                        wres *= self.w
                chi2 = wres.T.dot(wres)
                g = J.T.dot(wres)
                C = (J.T * J).todense()
                self.factor = scipy.linalg.cho_factor(C)
                deltax = scipy.linalg.cho_solve(self.factor, g).squeeze()
                chi2new = self(x + deltax, grad=False, hess=False, goods=goods)
                actual_red = chi2 - chi2new
                # print np.linalg.cond((J.T * J).todense())
                # For a linear f:
                pred_red = np.dot(deltax, g)
                rho = actual_red / pred_red
                dof = len(yi) - len(x)
                if verbose:
                    print(
                        (
                            "chi2=%.2f (dof=%.2f) -> %.2f (%.2f / %.2f = %g), |g|=%g"
                            % (
                                chi2,
                                dof,
                                chi2new,
                                actual_red,
                                pred_red,
                                rho,
                                np.sqrt(np.dot(g, g)),
                            )
                        )
                    )
                if (
                    not (fallback) or rho > 0
                ):  # accept the Gauss-Newton step even if it is rather poor
                    x += deltax
                    # Convergence criterion
                    if actual_red < 0.01:  # and pred_red < 0.01 and rho > 0.9:
                        break
                else:
                    print("bad step, giving up")
                    break
                it = it + 1
                if it > maxiter:
                    raise ValueError("Grosse fatigue")
                    if verbose:
                        print("Reach maximum number of iteration")
                    break
            chi2 = self(x, goods=goods)
            dof = len(yi) - len(x)
            if verbose:
                print(
                    (
                        "Converged in %d iterations: chi2=%g, d.o.f.=%d, chi2/dof=%.3f"
                        % (it, chi2, dof, chi2 / dof)
                    )
                )
            if keys.pop("full_output", False):
                return x, chi2, dof, self.factor
            else:
                return x
        else:
            raise ValueError("Unknown algorithm")


# def lsqfit_gauss_newton(y, model, p0, eq_constraints=None, weights=None, **keys):
#    """ Gauss-Newton with a confidence region check to switch to
#    steepest decent in case of trouble. Tipycally fails to
#    converge in high-residual cases. Try Newton or
#    quasi-Newton in this cases.
#    """
#    raise_exception = keys.pop('raise', False)
#    x = p0.copy()
#    while True:
#        m, J = model(0, x, jac=True)
#        if weights is not None:
#            J.data *= weights[J.row]
#        J = J.tocsr()
#        wres = y - m
#        if weights is not None:
#            wres *= weights
#        chi2 = wres.T.dot(wres)
#        g = J.T.dot(wres)
#        if eq_constraints is not None:
#            c, A = eq_constraints(x, jac=True)
#            H = scipy.sparse.bmat([[J.T*J, A],[A.T, None]])
#            B = np.hstack([g, -c])
#            lu = sparse.linalg.splu(H)
#            dxl = lu.solve(B)
#            dx = dxl[:-len(c)]
#            lagrange = dxl[-len(c):]
#            if it == 0:
#                    self.factor = cholesky_AAt(J.T)
#                else:
#                    self.factor.cholesky_AAt_inplace(J.T)
#                deltax = self.factor(g).squeeze()
#                chi2new = self(x + deltax, grad=False, hess=False)
#                actual_red = chi2 - chi2new
#                #print np.linalg.cond((J.T * J).todense())
#                # For a linear f:
#                pred_red = (np.dot(deltax, g))
#                rho = actual_red / pred_red
#                dof = len(self.yi) - len(x)
#                print "%.2f (%.2f) -> %.2f (%.2f / %.2f = %g), |g|=%g" % (chi2, dof, chi2new, actual_red, pred_red, rho, np.sqrt(np.dot(g,g)))
#                if rho > 0: # accept the Gauss-Newton step even if it is rather poor
#                    x += deltax
#                    # Convergence criterion
#                    if actual_red < 0.01 and pred_red < 0.01 and rho > 0.9:
#                        break
#                else:  # attempt to save the meuble
#                    print "bad step: trying steepest decent"
#                    def c(mu):
#                        return self(x + mu * g, grad=False, hess=False)
#                    mu, chi2new, niter, funcalls = scipy.optimize.brent(c, brack=(0,1), tol=1e-2, full_output=True)
#                    print "Brent found new minimum in %d iter: x + %e * g -> %.2f" % (niter, mu, chi2new)
#                    if mu < 0:
#                        if raise_exception:
#                            mu = np.linspace(-1e-4,1e-4,20)
#                            c = [self(x + _mu * g, grad=False, hess=False) for _mu in mu]
#                            import matplotlib.pyplot as plt
#                            plt.plot(mu, c)
#                            plt.show()
#                            raise ValueError("Dear Lord, help us out of this merdier")
#                        else:
#                            pass
#                    x += mu * g
#                    """
#                    print "bad step"
#                    d = np.bincount(J.indices, J.data * J.data)
#                    #deltax = g / d
#                    H = J.T * J
#                    d = scipy.sparse.diags(d, 0).tocsc()
#                    nu = 0.1
#                    while (chi2new > chi2):
#                        self.factor.cholesky_inplace(J.T*J + nu * d)
#                        #self.factor.cholesky_AAt_inplace(J.T, nu)
#                        deltax = self.factor(g).squeeze()
#                        chi2new = self(x + deltax, grad=False, hess=False)
#                        print chi2new
#                        nu *= 2
#                        if nu > 1e2:
#                            # Holly shit we are in a bad shape
#                            mu = np.linspace(0,1e-4,10)
#                            c = [self(x + _mu * g, grad=False, hess=False) for _mu in mu]
#                            import matplotlib.pyplot as plt
#                            plt.plot(mu, c)
#                            plt.show()
#                            raise ValueError("Dear Lord, help us out of this merdier")
#                            #x2 = scipy.optimize.anneal(self, x)
#                            #deltax = x2 - x
#                    print "smoothing param: %g (chi2: %.2f -> %.2f)" % (nu, chi2, chi2new)
#                    x += deltax
#                    #x = scipy.optimize.fmin(self, x)
#                    """
#                it = it + 1
#                if it > 2e3:
#                    raise ValueError("Grosse fatigue")
#            chi2 = self(x)
#            dof = len(self.yi) - len(x)
#            print "Converged in %d iterations: chi2=%g, d.o.f.=%d, chi2/dof=%.3f" % (it, chi2, dof, chi2/dof)
#            if keys.pop('full_output', False):
#                return x, chi2, dof, self.factor
#            else:
#                return x
#
def fmin_ncg(
    f,
    x0,
    fprime,
    fhess_p=None,
    fhess=None,
    args=(),
    avextol=1e-5,
    epsilon=_epsilon,
    maxiter=None,
    full_output=0,
    disp=1,
    retall=0,
    callback=None,
):
    """
    Unconstrained minimization of a function using the Newton-CG method.

    This is a patched version of the scipy.optimize function accepting
    sparse Hessian matrices.

    Parameters
    ----------
    f : callable ``f(x, *args)``
        Objective function to be minimized.
    x0 : ndarray
        Initial guess.
    fprime : callable ``f'(x, *args)``
        Gradient of f.
    fhess_p : callable ``fhess_p(x, p, *args)``, optional
        Function which computes the Hessian of f times an
        arbitrary vector, p.
    fhess : callable ``fhess(x, *args)``, optional
        Function to compute the Hessian matrix of f.
    args : tuple, optional
        Extra arguments passed to f, fprime, fhess_p, and fhess
        (the same set of extra arguments is supplied to all of
        these functions).
    epsilon : float or ndarray, optional
        If fhess is approximated, use this value for the step size.
    callback : callable, optional
        An optional user-supplied function which is called after
        each iteration.  Called as callback(xk), where xk is the
        current parameter vector.
    avextol : float, optional
        Convergence is assumed when the average relative error in
        the minimizer falls below this amount.
    maxiter : int, optional
        Maximum number of iterations to perform.
    full_output : bool, optional
        If True, return the optional outputs.
    disp : bool, optional
        If True, print convergence message.
    retall : bool, optional
        If True, return a list of results at each iteration.

    Returns
    -------
    xopt : ndarray
        Parameters which minimize f, i.e. ``f(xopt) == fopt``.
    fopt : float
        Value of the function at xopt, i.e. ``fopt = f(xopt)``.
    fcalls : int
        Number of function calls made.
    gcalls : int
        Number of gradient calls made.
    hcalls : int
        Number of hessian calls made.
    warnflag : int
        Warnings generated by the algorithm.
        1 : Maximum number of iterations exceeded.
    allvecs : list
        The result at each iteration, if retall is True (see below).

    See also
    --------
    minimize: Interface to minimization algorithms for multivariate
        functions. See the 'Newton-CG' `method` in particular.

    Notes
    -----
    Only one of `fhess_p` or `fhess` need to be given.  If `fhess`
    is provided, then `fhess_p` will be ignored.  If neither `fhess`
    nor `fhess_p` is provided, then the hessian product will be
    approximated using finite differences on `fprime`. `fhess_p`
    must compute the hessian times an arbitrary vector. If it is not
    given, finite-differences on `fprime` are used to compute
    it.

    Newton-CG methods are also called truncated Newton methods. This
    function differs from scipy.optimize.fmin_tnc because

    1. scipy.optimize.fmin_ncg is written purely in python using numpy
        and scipy while scipy.optimize.fmin_tnc calls a C function.
    2. scipy.optimize.fmin_ncg is only for unconstrained minimization
        while scipy.optimize.fmin_tnc is for unconstrained minimization
        or box constrained minimization. (Box constraints give
        lower and upper bounds for each variable separately.)

    References
    ----------
    Wright & Nocedal, 'Numerical Optimization', 1999, pg. 140.

    """
    opts = {
        "xtol": avextol,
        "eps": epsilon,
        "maxiter": maxiter,
        "disp": disp,
        "return_all": retall,
    }

    res = _minimize_newtoncg(
        f, x0, args, fprime, fhess, fhess_p, callback=callback, **opts
    )

    if full_output:
        retlist = (
            res["x"],
            res["fun"],
            res["nfev"],
            res["njev"],
            res["nhev"],
            res["status"],
        )
        if retall:
            retlist += (res["allvecs"],)
        return retlist
    else:
        if retall:
            return res["x"], res["allvecs"]
        else:
            return res["x"]


def _minimize_newtoncg(
    fun,
    x0,
    args=(),
    jac=None,
    hess=None,
    hessp=None,
    callback=None,
    xtol=1e-5,
    eps=_epsilon,
    maxiter=None,
    disp=False,
    return_all=False,
    **unknown_options,
):
    """
    Minimization of scalar function of one or more variables using the
    Newton-CG algorithm.

    This is a patched version of the scipy.optimize function accepting
    sparse hessian.

    Options for the Newton-CG algorithm are:
        disp : bool
            Set to True to print convergence messages.
        xtol : float
            Average relative error in solution `xopt` acceptable for
            convergence.
        maxiter : int
            Maximum number of iterations to perform.
        eps : float or ndarray
            If `jac` is approximated, use this value for the step size.

    This function is called by the `minimize` function with
    `method=Newton-CG`. It is not supposed to be called directly.

    Also note that the `jac` parameter (Jacobian) is required.
    """
    _check_unknown_options(unknown_options)
    if jac is None:
        raise ValueError("Jacobian is required for Newton-CG method")
    f = fun
    fprime = jac
    fhess_p = hessp
    fhess = hess
    avextol = xtol
    epsilon = eps
    retall = return_all

    x0 = np.asarray(x0).flatten()
    fcalls, f = _wrap_function(f, args)
    gcalls, fprime = _wrap_function(fprime, args)
    hcalls = 0
    if maxiter is None:
        maxiter = len(x0) * 200

    xtol = len(x0) * avextol
    update = [2 * xtol]
    xk = x0
    if retall:
        allvecs = [xk]
    k = 0
    old_fval = f(x0)
    old_old_fval = None
    float64eps = np.finfo(np.float64).eps
    warnflag = 0
    while (np.add.reduce(np.abs(update)) > xtol) and (k < maxiter):
        # Compute a search direction pk by applying the CG method to
        #  del2 f(xk) p = - grad f(xk) starting from 0.

        # b = -fprime(xk)
        c2, b, A = fun(xk, grad=True, hess=True)
        b = -b
        hcalls = hcalls + 1
        # gcalls = gcalls + 1
        maggrad = np.add.reduce(np.abs(b))
        eta = np.min([0.5, np.sqrt(maggrad)])
        termcond = eta * maggrad
        xsupi = np.zeros(len(x0), dtype=x0.dtype)
        ri = -b
        psupi = -ri
        i = 0
        dri0 = np.dot(ri, ri)

        # if fhess is not None:             # you want to compute hessian once.
        #    A = fhess(*(xk,) + args)
        #    hcalls = hcalls + 1

        while np.add.reduce(np.abs(ri)) > termcond:
            # if fhess is None:
            #    if fhess_p is None:
            #        Ap = approx_fhess_p(xk, psupi, fprime, epsilon)
            #    else:
            #        Ap = fhess_p(xk, psupi, *args)
            #        hcalls = hcalls + 1
            # else:
            # Ap = np.dot(A, psupi)
            Ap = A.dot(psupi)
            # check curvature
            Ap = np.asarray(Ap).squeeze()  # get rid of matrices...
            curv = np.dot(psupi, Ap)
            if 0 <= curv <= 3 * float64eps:
                break
            elif curv < 0:
                if i > 0:
                    break
                else:
                    # fall back to steepest descent direction
                    xsupi = dri0 / (-curv) * b
                    break
            alphai = dri0 / curv
            xsupi = xsupi + alphai * psupi
            ri = ri + alphai * Ap
            dri1 = np.dot(ri, ri)
            betai = dri1 / dri0
            psupi = -ri + betai * psupi
            i = i + 1
            dri0 = dri1  # update np.dot(ri,ri) for next time.

        pk = xsupi  # search direction is solution to system.
        gfk = -b  # gradient at xk

        try:
            alphak, fc, gc, old_fval, old_old_fval, gfkp1 = _line_search_wolfe12(
                f, fprime, xk, pk, gfk, old_fval, old_old_fval
            )
        except _LineSearchError:
            # Line search failed to find a better solution.
            warnflag = 2
            break

        update = alphak * pk
        xk = xk + update  # upcast if necessary
        if callback is not None:
            callback(xk)
        if retall:
            allvecs.append(xk)
        k += 1

    fval = old_fval
    if warnflag == 2:
        msg = _status_message["pr_loss"]
        if disp:
            print(("Warning: " + msg))
            print(("         Current function value: %f" % fval))
            print(("         Iterations: %d" % k))
            print(("         Function evaluations: %d" % fcalls[0]))
            print(("         Gradient evaluations: %d" % gcalls[0]))
            print(("         Hessian evaluations: %d" % hcalls))
    elif k >= maxiter:
        warnflag = 1
        msg = _status_message["maxiter"]
        if disp:
            print(("Warning: " + msg))
            print(("         Current function value: %f" % fval))
            print(("         Iterations: %d" % k))
            print(("         Function evaluations: %d" % fcalls[0]))
            print(("         Gradient evaluations: %d" % gcalls[0]))
            print(("         Hessian evaluations: %d" % hcalls))
    else:
        msg = _status_message["success"]
        if disp:
            print(msg)
            print(("         Current function value: %f" % fval))
            print(("         Iterations: %d" % k))
            print(("         Function evaluations: %d" % fcalls[0]))
            print(("         Gradient evaluations: %d" % gcalls[0]))
            print(("         Hessian evaluations: %d" % hcalls))

    result = OptimizeResult(
        fun=fval,
        jac=gfk,
        nfev=fcalls[0],
        njev=gcalls[0],
        nhev=hcalls,
        status=warnflag,
        success=(warnflag == 0),
        message=msg,
        x=xk,
    )
    if retall:
        result["allvecs"] = allvecs
    return result


class RobustMinimizer:
    """ """

    def __init__(self, model, y, x=None, weights=None, verbose=1):
        """ """
        self.model = model
        self.x = x
        self.y = y
        self.w = weights
        self.verbose = verbose
        self._chi2 = Chi2(model, x, y, w=weights)
        N = len(self.y)
        self.bads = np.zeros(N).astype(bool)
        self.suppress = None
        self.restore = None

    def minimize(self, p0, algorithm="newton_homebrew", bads=None, **kwargs):
        """ """
        # restore the initial bads
        if bads is None:
            self.bads = np.zeros(len(self.y)).astype(bool)
        else:
            self.bads = bads

        # first minimization
        p = self._chi2.minimize(p0, goods=(~self.bads), algorithm=algorithm, **kwargs)

        niter = 1
        while self.flag_outliers(p, **kwargs):
            logging.info("nbads/ntot=%d/%d" % (self.bads.sum(), len(self.bads)))

            pp = self._chi2.minimize(
                p, goods=(~self.bads), algorithm=algorithm, **kwargs
            )
            niter += 1
            logging.info(
                "dp=%g chi2=%g"
                % (np.sum(np.abs(pp - p)), self._chi2(pp, goods=~self.bads))
            )
            p[:] = pp

        if self.verbose:
            logging.info(
                "converged after %d iterations %d/%d outliers removed"
                % (niter, self.bads.sum(), len(self.bads))
            )

        self.params = self.model.params.copy()
        self.params.free = p
        self.eparams = self.model.params.copy()
        self.eparams.free = np.sqrt(self.get_diag_cov(p))

        return p

    def check_local_param(self, p, local_param):
        """ """
        v, J = self.model(self.x, p, jac=True)
        local_cols = slice(
            self.model.params[local_param].indexof(0),
            self.model.params[local_param].indexof(-1),
        )
        J = J.tocsc()[~self.bads, local_cols]
        return J.sum(axis=0)

    def suppress_one_at_a_time(self, r, p, local_param):
        """
        For each paramter, suppress only the worst outlier

        """
        if self.suppress.sum() == 0:
            return self.suppress

        v, J = self.model(self.x, p, jac=1)
        local_cols = slice(
            self.model.params[local_param].indexof(0),
            self.model.params[local_param].indexof(-1),
        )
        J = J.tocsc()[self.suppress, local_cols].tocoo()
        J.data = r[self.suppress][J.row]

        # the line number of the largest outlier
        i_max = np.array(J.argmax(axis=0)).squeeze()
        # it may also be that the corresponding column is empty
        v_max = np.array(J.max(axis=0).todense()).squeeze()
        # the outliers we need to remove
        to_remove = i_max[i_max > 0 | ((i_max == 0) & (v_max > 0.0))]

        suppress_one = np.zeros_like(self.suppress[self.suppress])
        suppress_one[to_remove] = True
        n_before = self.suppress.sum()
        self.suppress[self.suppress] = suppress_one
        n_after = self.suppress.sum()
        logging.info("to suppress: %d -> %d" % (n_before, n_after))

    def flag_outliers(self, p, nsig=5, local_param=None):
        """flag the outliers that deviate from the mean chi2 by at least
        nsig. The standard deviation of the residuals is the square root of the chi2
        divided by the number of degrees of freedom.

        Args:
           p (array of floats): current solution (free parameters)
           nsig (int or float): cut in sigma
           local_param (array of ints): parameters

        Returns:
           None

        .. note: modifies self.bads, self.suppress, self.restore
        """
        r = self._chi2.get_res(p)
        r *= r
        if self.bads is not None:
            chi2 = self._chi2(p, goods=~self.bads)
            nmeas = (~self.bads).sum()
        else:
            chi2 = self._chi2(p)
            nmeas = len(r)

        n_free_pars = len(p)
        ndof = nmeas - n_free_pars

        # identify the outliers
        bads = r > (nsig**2 * chi2 / ndof)

        # now, who should I actually suppress & restore ?
        self.suppress = bads & ~self.bads
        self.restore = ~bads & self.bads

        # here, I should add the local params rejection
        if local_param is not None:
            self.suppress_one_at_a_time(r, p, local_param)

        n_suppressed = self.suppress.sum()
        n_restored = self.restore.sum()

        if n_suppressed == 0:
            self.suppress = None
        else:
            self.bads[self.suppress] = True

        if n_restored == 0:
            self.restore = None
        else:
            self.bads[self.restore] = False

        # update the chi2
        new_chi2 = self._chi2(p, goods=~self.bads)

        if self.verbose > 0:
            logging.info("suppressed %d outliers" % n_suppressed)
            logging.info("restored %d outliers" % n_restored)
            logging.info("chi2: %g -> %g" % (chi2, new_chi2))

        return n_restored > 0 or n_suppressed > 0

    def get_res(self, p=None, x=None):
        """
        how do we deal with the bads ?
        """
        if p is None:
            pp = self.model.params.free
        else:
            pp = p
        r = self.model(x, p=pp)
        r -= self.y
        return r

    def get_wres(self, x=None):
        """
        how do we deal with the bads ?
        """
        r = self.model(x, p=self.params.free)
        r -= self.y
        return self.weights * r

    def get_cov(self):
        return self._chi2.get_cov(self.params.free, goods=~self.bads)

    def get_block_cov(self):
        pass

    def get_diag_cov(self, p):
        return self._chi2.get_diag_cov(p, goods=~self.bads)

    def ndof(self):
        nmeas = (~self.bads).sum()
        npars = len(self.params.free)
        return nmeas - npars

    def chi2(self, p=None):
        if p:
            return self._chi2(p, goods=~self.bads)
        return self._chi2(self.params.free, goods=~self.bads)


class OptimizeResult(dict):
    """Represents the optimization result.

    Attributes
    ----------
    x : ndarray
        The solution of the optimization.
    success : bool
        Whether or not the optimizer exited successfully.
    status : int
        Termination status of the optimizer. Its value depends on the
        underlying solver. Refer to `message` for details.
    message : str
        Description of the cause of the termination.
    fun, jac, hess, hess_inv : ndarray
        Values of objective function, Jacobian, Hessian or its inverse (if
        available). The Hessians may be approximations, see the documentation
        of the function in question.
    nfev, njev, nhev : int
        Number of evaluations of the objective functions and of its
        Jacobian and Hessian.
    nit : int
        Number of iterations performed by the optimizer.
    maxcv : float
        The maximum constraint violation.

    Notes
    -----
    There may be additional attributes not listed above depending of the
    specific solver. Since this class is essentially a subclass of dict
    with attribute accessors, one can see which attributes are available
    using the `keys()` method.
    """

    def __getattr__(self, name):
        try:
            return self[name]
        except KeyError:
            raise AttributeError(name)

    __setattr__ = dict.__setitem__
    __delattr__ = dict.__delitem__

    def __repr__(self):
        if list(self.keys()):
            m = max(list(map(len, list(self.keys())))) + 1
            return "\n".join(
                [k.rjust(m) + ": " + repr(v) for k, v in list(self.items())]
            )
        else:
            return self.__class__.__name__ + "()"
