# This file is part of the pyMOR project (http://www.pymor.org).
# Copyright Holders: Felix Albrecht, Rene Milk, Stephan Rave
# License: BSD 2-Clause License (http://opensource.org/licenses/BSD-2-Clause)

from __future__ import absolute_import, division, print_function

from pymor import defaults
from pymor.core import getLogger
from pymor.core.exceptions import InversionError, NewtonError


def newton(operator, rhs, initial_guess=None, mu=None, error_norm=None, maxiter=None, reduction=None, abs_limit=None):
    maxiter = defaults.newton_maxiter if maxiter is None else maxiter
    reduction = defaults.newton_reduction if reduction is None else reduction
    abs_limit = defaults.newton_abs_limit if abs_limit is None else abs_limit
    logger = getLogger('pymor.algorithms.newton')
    if initial_guess is None:
        initial_guess = operator.type_source.zeros(operator.dim_source)

    U = initial_guess.copy()
    residual = rhs - operator.apply(U, mu=mu)
    err = first_err = residual.l2_norm()[0] if error_norm is None else error_norm(residual)[0]
    logger.info('      Initial Resiudal: {:5e}'.format(err))

    iteration = 0
    while iteration < maxiter and err > abs_limit and err/first_err > reduction:
        iteration += 1
        jacobian = operator.jacobian(U, mu=mu)
        try:
            correction = jacobian.apply_inverse(residual)
        except InversionError:
            raise NewtonError('Could not invert jacobian')
        U += correction
        residual = rhs - operator.apply(U, mu=mu)
        old_err = err
        err = residual.l2_norm()[0] if error_norm is None else error_norm(residual)[0]
        logger.info('Iteration {:2}: Residual: {:5e},  Reduction: {:5e}, Total Reduction: {:5e}'
                    .format(iteration, err, err / old_err, err / first_err))

    if err > abs_limit and err/first_err > reduction:
        raise NewtonError('Failed to converge')

    return U

@@ -459,16 +459,27 @@
 rm=None,
+%0A          
  maxiter
@@ -515,16 +515,72 @@
 mit=None
+,%0A           return_stages=False, return_residuals=False
 ):%0A    m
@@ -846,24 +846,40 @@
 ms.newton')%0A
+%0A    data = %7B%7D%0A%0A
     if initi
@@ -969,16 +969,212 @@
 ource)%0A%0A
+    if return_stages:%0A        data%5B'stages'%5D = operator.type_source.empty(operator.dim_source)%0A%0A    if return_residuals:%0A        data%5B'residuals'%5D = operator.type_range.empty(operator.dim_range)%0A%0A
     U = 
@@ -1236,24 +1236,25 @@
 y(U, mu=mu)%0A
+%0A
     err = fi
@@ -1488,32 +1488,211 @@
 err %3E reduction:
+  # or iteration == 0:%0A        if iteration %3E 0 and return_stages:%0A            data%5B'stages'%5D.append(U)%0A        if return_residuals:%0A            data%5B'residuals'%5D.append(residual)
 %0A        iterati
@@ -1973,32 +1973,33 @@
 apply(U, mu=mu)%0A
+%0A
         old_err 
@@ -2365,24 +2365,24 @@
 converge')%0A%0A
-
     return U
@@ -2381,9 +2381,15 @@
 return U
+, data
 %0A

add options to return intermediate stages and residuals