import math
import matplotlib.pyplot as plt
import pacopy


class SimpleScalarProblem:
    def inner(self, a, b):
        """The inner product in the problem domain. For scalars, this is just a
        multiplication.
        """
        return a * b

    def norm2_r(self, a):
        """The norm in the range space; used to determine if a solution has been found."""
        return a**2

    def f(self, u, lmbda):
        """The evaluation of the function to be solved"""
        return math.sin(u) - lmbda

    def df_dlmbda(self, u, lmbda):
        """The function's derivative with respect to the parameter. Used in Euler-Newton
        continuation.
        """
        return -1.0

    def jacobian_solver(self, u, lmbda, rhs):
        """A solver for the Jacobian problem. For scalars, this is just a division."""
        return rhs / math.cos(u)


problem = SimpleScalarProblem()

lmbda_list = []
values_list = []


def callback(k, lmbda, sol):
    # Use the callback for plotting, writing data to files etc.
    lmbda_list.append(lmbda)
    values_list.append(sol)


pacopy.euler_newton(
    problem,
    u0=0.0,
    lmbda0=0.0,
    callback=callback,
    max_steps=20,
    newton_tol=1.0e-10,
    verbose=False,
)

# plot solution
plt.plot(values_list, lmbda_list, ".-")
plt.xlabel("$u_1$")
plt.ylabel("$\\lambda$")
plt.show()

import matplotlib.pyplot as plt
import numpy as np
import pacopy
import scipy
from mpl_toolkits.mplot3d import Axes3D


class SimpleProblem2D:
    def inner(self, a, b):
        return np.dot(a, b)

    def norm2_r(self, a):
        return np.dot(a, a)

    def f(self, u, lmbda):
        return np.array(
            [
                np.sin(u[0]) - lmbda - u[1] ** 2,
                np.cos(u[1]) * u[1] - lmbda,
            ]
        )

    def df_dlmbda(self, u, lmbda):
        return -np.ones_like(u)

    def jacobian_solver(self, u, lmbda, rhs):
        A = np.array(
            [
                [np.cos(u[0]), -2 * u[1]],
                [0.0, np.cos(u[1]) - np.sin(u[1]) * u[1]],
            ]
        )
        return np.linalg.solve(A, rhs)


problem = SimpleProblem2D()
# Initial guess and initial parameter value
u0 = np.zeros(2)
lmbda0 = 0.0

# Init lists
lmbda_list = []
values_list = []


def callback(k, lmbda, sol):
    lmbda_list.append(lmbda)
    values_list.append(sol)


pacopy.euler_newton(problem, u0, lmbda0, callback, max_steps=50, newton_tol=1.0e-10)

# plot solution
fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
ax.plot(*np.array(values_list).T, lmbda_list, ".-")
ax.set_xlabel("$u_1$")
ax.set_ylabel("$u_2$")
ax.set_zlabel("$\\lambda$")
# plt.show()
plt.savefig("simple2d.svg", transparent=True, bbox_inches="tight")
plt.close()

import matplotlib.pyplot as plt
import numpy as np
import pacopy
import scipy.sparse
import scipy.sparse.linalg


# This is the classical finite-difference approximation
class Bratu1d:
    def __init__(self, n):
        self.n = n
        h = 1.0 / (self.n - 1)

        self.H = np.full(self.n, h)
        self.H[0] = h / 2
        self.H[-1] = h / 2

        self.A = (
            scipy.sparse.diags([-1.0, 2.0, -1.0], [-1, 0, 1], shape=(self.n, self.n))
            / h**2
        )

    def inner(self, a, b):
        """The inner product of the problem. Can be np.dot(a, b), but factoring in
        the mesh width stays true to the PDE.
        """
        return np.dot(a, self.H * b)

    def norm2_r(self, a):
        """The norm in the range space; used to determine if a solution has been found."""
        return np.dot(a, a)

    def f(self, u, lmbda):
        """The evaluation of the function to be solved"""
        out = self.A.dot(u) - lmbda * np.exp(u)
        out[0] = u[0]
        out[-1] = u[-1]
        return out

    def df_dlmbda(self, u, lmbda):
        """The function's derivative with respect to the parameter. Used in Euler-Newton
        continuation.
        """
        out = -np.exp(u)
        out[0] = 0.0
        out[-1] = 0.0
        return out

    def jacobian_solver(self, u, lmbda, rhs):
        """A solver for the Jacobian problem."""
        M = self.A.copy()
        d = M.diagonal().copy()
        d -= lmbda * np.exp(u)
        M.setdiag(d)
        # Dirichlet conditions
        M.data[0][self.n - 2] = 0.0
        M.data[1][0] = 1.0
        M.data[1][self.n - 1] = 1.0
        M.data[2][1] = 0.0
        return scipy.sparse.linalg.spsolve(M.tocsr(), rhs)


problem = Bratu1d(51)
# Initial guess and parameter value
u0 = np.zeros(problem.n)
lmbda0 = 0.0

lmbda_list = []
values_list = []


def callback(k, lmbda, sol):
    # Use the callback for plotting, writing data to files etc.
    fig = plt.figure()
    ax1 = fig.add_subplot(111)
    ax1.set_xlabel("$\\lambda$")
    ax1.set_ylabel("$||u||_2$")
    ax1.grid()

    lmbda_list.append(lmbda)
    # use the norm of the currentsolution for plotting on the y-axis
    values_list.append(np.sqrt(problem.inner(sol, sol)))

    ax1.plot(lmbda_list, values_list, "-x", color="C0")
    ax1.set_xlim(0.0, 4.0)
    ax1.set_ylim(0.0, 6.0)
    plt.close()


# Natural parameter continuation
# pacopy.natural(problem, u0, lmbda0, callback, max_steps=100)

pacopy.euler_newton(
    problem, u0, lmbda0, callback, max_steps=500, max_num_retries=10, newton_tol=1.0e-10
)