# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_poisson_solvers.ipynb.

# %% auto 0
__all__ = ['poisson_gauss_seidel_with_sor_solver', 'poisson_non_iterative_solver', 'poisson_newton_solver',
           'poisson_newton_alternative_solver']

# %% ../nbs/01_poisson_solvers.ipynb 5
from . import core

from functools import partial

import numpy as np
import scipy
from scipy import sparse


# %% ../nbs/01_poisson_solvers.ipynb 9
def poisson_gauss_seidel_with_sor_solver(
    w, r=None, verbose=False, log_middle_values=False,
):
    N = w.shape[0] - 1
    nx = ny = N + 1

    h = 1 / N

    # Handle the SOR parameter
    if r is None:
        r = 2 / (1 + np.pi / N) # optimal value

    psi = np.zeros((nx, ny)) # streamfunction

    middle_values = []

    for iteration in range(1, 4 * nx + 1):
        for i in range(1, nx - 1):
            for j in range(1, ny - 1):
                psi[i, j] = (1 - r) * psi[i, j] + r * (
                    psi[i - 1, j] + psi[i + 1, j] +
                    psi[i, j - 1] + psi[i, j + 1] +
                    w[i, j] * h ** 2
                ) / 4
        
        if verbose:
            print(f"{iteration=}; psi(0.5, 0.5) = {psi[nx // 2, ny // 2]}")
        
        if log_middle_values:
            middle_values.append(psi[nx // 2, ny // 2])
    
    return psi, iteration, middle_values


# %% ../nbs/01_poisson_solvers.ipynb 24
def construct_laplacian_kernel_matrix(N, h=None):
    '''
    Construct the matrix that defines the linear system in sparse form
    '''

    if h is None:
        h = 1 / N
    
    # The output size is N - 1 because we are excluding the boundary points
    N -= 1

    values = np.ones(5 * N**2 - 4 * N) / h**2
    values[:N**2] *= -4

    all_indices = np.arange(N ** 2)
    
    rows = [np.arange(N ** 2)]
    cols = [np.arange(N ** 2)]

    rows.append(all_indices[N:N**2])
    cols.append(all_indices[:-N])

    rows.append(all_indices[:-N])
    cols.append(all_indices[N:N**2])

    rows.append(all_indices[all_indices % N != 0])
    cols.append(all_indices[(all_indices + 1) % N != 0])

    rows.append(all_indices[(all_indices + 1) % N != 0])
    cols.append(all_indices[all_indices % N != 0])

    # Flatten the lists
    rows = np.concatenate(rows)
    cols = np.concatenate(cols)

    # Create the sparse matrix from the above information
    return sparse.csr_matrix((values, (rows, cols)), shape=(N ** 2, N ** 2))


# %% ../nbs/01_poisson_solvers.ipynb 25
def poisson_non_iterative_solver(w, algorithm="bicgstab"):
    N = w.shape[0] - 1

    # Construct the matrix that defines the linear system
    kernel_matrix = construct_laplacian_kernel_matrix(N)
    
    # Cast vorticity to the required form
    w = w[1:-1, 1:-1].flatten()

    # Solve the linear system.
    # Note -ve sign is applied to make A positive definite which is required
    # by some of the iterative solvers.
    psi = core.solve_sparse_linear_system(
        A=-kernel_matrix, b=w, algorithm=algorithm
    )
    
    psi = psi.reshape(N-1, N-1) # exclude the boundary points

    psi = np.pad(psi, (1, 1), mode="constant")
    
    return psi


# %% ../nbs/01_poisson_solvers.ipynb 32
def get_jacobian(N, h=1):
    return construct_laplacian_kernel_matrix(N, h=h)
    

def f(x, w):
    N = int(np.sqrt(x.shape[0])) + 1
    h = 1 / N

    x = x.reshape(N-1, N-1)

    # note we compute h ** 2 * f_current
    f = -4 * x + h ** 2 * w[1:-1, 1:-1]
    f[1:, :] += x[:-1, :]
    f[:-1, :] += x[1:, :]
    f[:, :-1] += x[:, 1:]
    f[:, 1:] += x[:, :-1]

    return f.flatten()


def newton_iterator(
    w, f, get_jacobian,
    algorithm="bicgstab", TOL=1e-8, max_iter=10, quiet=True
):
    '''
        - w: vorticity
        - f: evaluates the function given x, w
        - get_jacobian: evaluates the Jacobian given N
    '''

    N = w.shape[0] - 1
    
    n_iter = 0 # number of iterations

    # Initialization
    x = np.zeros(((N - 1) ** 2))
    f_current = f(x, w)
    
    # Check if the initial guess is a solution
    f_norm = scipy.linalg.norm(f_current)
    if f_norm <= TOL:
        if not quiet:
            print(f"n_iter={n_iter}")

        return x, n_iter
    
    
    while n_iter < max_iter:
        n_iter += 1
        # Set h=1 to have better scaling and note that we actually compute h ** 2 * f_current
        jacobian = get_jacobian(N)
        
        dx = core.solve_sparse_linear_system(
            A=jacobian, b=-f_current, algorithm=algorithm
        )
        x_next = x + dx
        
        f_current = f(x_next, w)
        
        f_norm = scipy.linalg.norm(f_current)
        if not quiet:
            print(f"iter={n_iter}; |residual|={f_norm}; |dx|={scipy.linalg.norm(dx)}")
        if f_norm <= TOL:
            break
        
        x = x_next

    if not quiet:
        print(f"n_iter={n_iter}")
    
    return x_next, n_iter


# %% ../nbs/01_poisson_solvers.ipynb 33
def poisson_newton_solver(
    w, algorithm="bicgstab", TOL=1e-8, max_iter=10, quiet=True
):
    solution, _ = newton_iterator(
        w=w, f=f, get_jacobian=get_jacobian,
        algorithm=algorithm,
        TOL=TOL, max_iter=max_iter, quiet=quiet
    )
    
    N = w.shape[0] - 1
    psi = solution.reshape(N - 1, N - 1)

    psi = np.pad(psi, (1, 1), mode="constant", constant_values=0)
    
    return psi


# %% ../nbs/01_poisson_solvers.ipynb 34
def poisson_newton_alternative_solver(w, **kwargs):
    N = w.shape[0] - 1

    # Can define jacobian as `jac=lambda x, args: get_jacobian(N=N, h=1/N)`
    solution = scipy.optimize.root(
        fun=f, x0=np.zeros((N - 1) ** 2),
        args=(w, ),
        **kwargs
    )
    
    psi = solution.x.reshape(N - 1, N - 1)

    psi = np.pad(psi, (1, 1), mode="constant", constant_values=0)
    
    return psi, solution

