"""Solve a nonlinear diffusion problem via Newton-Raphson DG.

Problem: -d/dx((1 + u²) du/dx) = f  on [0, 1] with homogeneous Dirichlet BC.

Manufactured solution: u_exact(x) = sin(πx), u(0) = u(1) = 0.

RHS (computed via manufactured solution):
    f(x) = π² sin(πx) (3 sin²(πx) − 1)

The diffusion coefficient A(x, u) = (1 + u²) * I depends nonlinearly on u
but not on x.  It is supplied as a callable ``(x, u_val) -> matrix``:

    A = lambda x, u_val: (1.0 + u_val[0]**2) * jnp.eye(1)

The same callable is passed to NLSIPGFlux so that both volume and face
terms use the consistent A_eff = (1 + u²) * I.

Newton convergence: ``max_iter=20`` is required because the default (2)
is designed for linear PDEs.  With 20 iterations the solver converges to
machine precision for all mesh sizes used here.

Expected convergence rates (SIPG):  p=1 → L2 O(h²), H1 O(h);  p=2 → L2 O(h³), H1 O(h²).
"""

import jax.numpy as jnp
import matplotlib.pyplot as plt

from scimba_jax.linear_approximation.basis.analytic_bases import local_taylor_basis
from scimba_jax.linear_approximation.basis.general_bases import AnalyticBasis
from scimba_jax.linear_approximation.error_analysis import plot_convergence
from scimba_jax.linear_approximation.galerkin.dg.elliptic_dg_scheme import (
    EllipticDGscheme,
)
from scimba_jax.linear_approximation.galerkin.dg.flux import NLSIPGFlux
from scimba_jax.linear_approximation.meshes.mesh import Mesh
from scimba_jax.linear_approximation.quad.gauss_quad import UnitSquareTensorized
from scimba_jax.linear_approximation.variables.variables_dg import VariablesDG
from scimba_jax.mapping.mapping import InvertibleFunction, Mapping
from scimba_jax.physical_models.abstract_physical_weak_model import (
    AbstractPhysicalWeakModel,
)
from scimba_jax.physical_models.classical_weakform.nl_diffusion_advection_reaction_weak_form import (
    NLEllipticWeakForm,
)

physical_dim = 1
quad_order = 4
out_dim = 1

mapping_id = InvertibleFunction(lambda x: x, lambda y: y)

pde = NLEllipticWeakForm(
    dim=physical_dim,
    A=lambda x, u_val: (1.0 + u_val[0] ** 2) * jnp.eye(1),
    f=lambda x: jnp.pi**2
    * jnp.sin(jnp.pi * x[0])
    * (3.0 * jnp.sin(jnp.pi * x[0]) ** 2 - 1.0),
)


def u_exact_fn(x):
    return jnp.sin(jnp.pi * x[0])


# ── Convergence study (SIPG, degrees p = 1 and p = 2) ─────────────────────────
n_cells_list = [5, 10, 20, 40, 80]

fig, axes = plt.subplots(1, 3, figsize=(15, 4))
ax_l2, ax_h1, ax_linf = axes

for deg in [1, 2]:
    basis_order = deg

    def make_solver(b=basis_order, p=deg):
        def solver(n_cells):
            m = Mesh(
                dim=physical_dim,
                n_cells=[n_cells],
                ref_quad=UnitSquareTensorized(dim=physical_dim, order=quad_order),
                mapping=Mapping(mappings=[mapping_id]),
            )
            taylor_basis = AnalyticBasis(
                nb_basis=b + 1,
                out_dim=out_dim,
                mesh=m,
                local_basis=lambda coords, i, mesh, _b=b: local_taylor_basis(
                    coords, i, mesh, order=_b, out_dim=out_dim
                ),
                basis_type="scalar",
            )
            variables = VariablesDG(basis=taylor_basis, nb_variables=out_dim)

            def dirichlet_bc(__x):
                return jnp.zeros(out_dim)

            assembler = EllipticDGscheme(
                AbstractPhysicalWeakModel.from_weak_form(pde, dirichlet=dirichlet_bc),
                variables,
                NLSIPGFlux(sigma=p * (p + 1), h=1 / n_cells),
            )

            return EllipticDGscheme.solve(assembler, max_iter=20)

        return solver

    print(f"\nConvergence study SIPG p={deg}:")
    plot_convergence(
        solver=make_solver(),
        u_exact_fn=u_exact_fn,
        n_cells_list=n_cells_list,
        label=f"p={deg}",
        ax_l2=ax_l2,
        ax_h1=ax_h1,
        ax_linf=ax_linf,
        relative=True,
        title_sufix=" - SIPG",
    )

plt.suptitle(
    r"Nonlinear diffusion: $-\partial_x((1+u^2)\,\partial_x u) = f$", fontsize=13
)
plt.tight_layout()
plt.show()