# coding=utf-8

import numpy as np
import sys

from . import series_expansion as se

class ConvergenceError(Exception):
    pass

class MaxIterReached(ConvergenceError):
    pass
class DivergenceException(ConvergenceError):
    pass

def _nonlin_solve_1d_iter(prev, grid, basis_vals, Fs, Fs_args):
    order = grid.shape[0]
    N     = len(Fs)

    # evaluate the previous iteration on the grid
    prev_vals = []
    for diff_order in range(N - 1):
        prev_vals.append(prev.eval(grid, diff_order))

    # evaluate the RHS functions using the previous iteration
    F_vals = []
    for F in Fs:
        args = [grid, prev_vals]
        if Fs_args is not None:
            args += Fs_args
        F_vals.append(F(*args))

    # TODO should be doable with fewer explicit loops
    mat = np.copy(basis_vals[-1])
    for idx in range(order):
        for diff_order in range(N - 1):
            mat[:, idx] -= basis_vals[diff_order][:, idx] * F_vals[diff_order + 1]

    rhs = F_vals[0] - prev.eval(grid, N - 1)

    return np.linalg.solve(mat, rhs)

def nonlin_solve_1d(initial_guess, Fs, args = None, maxiter = 100, atol = 1e-14, grid = None, verbose = True):
    """
    Solve a non-linear ODE using a spectral method with Newton iteration.

    The ODE is expected to be of the form
        d^n u
        -----  = F(x, u(x), u'(x), u''(x), ... , u^(n-1)(x))
        dx^n

    Args:
        initial_guess (SeriesExpansion): The initial guess, from which the iteration is started.
                                         The spectral basis and its order are inferred from the
                                         initial guess.
        Fs (list of callables): The right-hand side of the ODE and its variational derivatives. I.e. Fs[0] is F,
                                Fs[i] is δF / δu^(i). The order of the equation n is determined by
                                the length of Fs minus one. Each callable must accept 3 parameters:
                                    - a numpy array containing the values of x
                                    - an iterable of n numpy arrays containing the values of u and
                                      its derivatives at x
                                    - a tuple of additional parameters passed to nonlin_solve_1d as args
        args (tuple or None): If not None, extra parameters passed to the Fs.
        maxiter: Maximum number of iterations done by the solver.
        atol: the difference between the coefficients in subsequent iterations below which convergence
              is assumed.
        grid: If non-None, the collocation grid to use. If left as None, the default grid from the basis
              object is used.
    Returns:
        The solution to the equation (a SeriesExpansion object).
    Raises:
        MaxIterReached: If the iteration fails to converge.
    """
    N = len(Fs)

    order = len(initial_guess.coeffs)
    if grid is None:
        grid = initial_guess.basis[0].colloc_grid(order)

    # precompute the values of the basis functions and their derivatives
    # at the grid
    basis = initial_guess.basis[0]
    basis_vals = []
    for diff_order in range(N):
        basis_val = np.empty((order, order))
        for idx in range(order):
            basis_val[:, idx] = basis.eval(idx, grid, diff_order)
        basis_vals.append(basis_val)

    solution = initial_guess
    for i in range(maxiter):
        delta = _nonlin_solve_1d_iter(solution, grid, basis_vals, Fs, args)

        err = np.max(np.abs(delta))
        if not np.isfinite(err):
            raise DivergenceException('nan')
        if err < atol:
            return solution
        solution = se.SeriesExpansion(solution.coeffs + delta, solution.basis)
        if verbose:
            sys.stderr.write('delta: %g, coeffs[0]: %g\n' % (err, solution.coeffs[0]))

    raise MaxIterReached('The horizon finder failed to converge')

def nonlin_residual(solution, N, grid, F, args):
    """
    Calculate the residual for the solution computed by nonlin_solve_1d() using
    2nd order finite differences.

    Args:
        solution (SeriesExpansion): The solution, as returned by nonlin_solve_1d().
        N (int): Order of the equation.
        grid (ndarray): An equidistant array of coordinates on which to evaluate the solution.
        F: The right-hand side of the ODE.
        args: Extra args to pass to F.
    """
    sol_vals = [solution.eval(grid)]
    dx = abs(grid[1] - grid[0])
    for diff_order in range(N):
        sol_vals.append(np.gradient(sol_vals[-1], dx))
    rhs = F(grid, sol_vals[:-1], args)
    return sol_vals[-1] - rhs