"""CG-FEM on the unit square with Robin conditions.

    -Delta u = 0                 in (0,1)^2
     grad u . n + alpha u = g    on the Robin sides

with the exact solution ``u(x, y) = 1 - x/2``, chosen so that it does **not**
vanish on the Robin boundary: ``u = 1/2`` on the east side, ``u = 1 - x/2`` on
south and north. That matters. With ``u = 0`` there, the ``alpha u`` term would
be zero and the test would pass just as well with ``alpha`` dropped entirely --
it would only be checking Neumann. Here dropping it changes ``g`` by
``alpha u != 0`` and the solution is visibly wrong.

Robin is the only condition whose prescribed flux depends on the solution, so
it is also the only one contributing to the *Jacobian* of the boundary term,
not just to the residual. Two configurations exercise that:

- **mixed**: Dirichlet west, Robin east, Neumann south/north;
- **pure Robin**: all four sides, so there is no constrained DOF anywhere and
  the whole problem is held together by the boundary integral. Coercivity comes
  from ``alpha > 0`` alone; if the ``alpha u`` Jacobian term were missing, the
  operator would be singular rather than merely inaccurate.

Being affine, the solution lies in every polynomial space used here, so the
error measures the *scheme*, not the approximation.
"""

# %%
import jax.numpy as jnp

from scimba_jax.linear_approximation.basis.analytic_bases import local_lagrange_basis
from scimba_jax.linear_approximation.basis.general_bases import AnalyticBasis
from scimba_jax.linear_approximation.error_analysis import l2_error
from scimba_jax.linear_approximation.galerkin.fem.elliptic_fe_scheme import (
    EllipticFEscheme,
)
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_fe import VariablesFE
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.laplacian_weak_form import (
    LaplacianWeakForm,
)
from scimba_jax.physical_models.weak_boundary_conditions import (
    Dirichlet,
    Neumann,
    Robin,
)
from scimba_jax.plots.plots_galerkin import plot_boundary_profiles, plot_solution_2d

DIM = 2
POLY_ORDER = 1
QUAD_ORDER = 3
N_CELLS = 8
OUT_DIM = 1
ALPHA = 2.0


def u_exact(x):
    """Exact solution ``u(x, y) = 1 - x/2``."""
    return jnp.array([1.0 - 0.5 * x[0]])


# grad u = (-1/2, 0), so grad u . n is -1/2 on east, +1/2 on west, 0 elsewhere.
# The Robin data is then g = grad u . n + alpha u, evaluated on each side.
def g_east(x):
    """Robin data on the east side, where ``u = 1/2``."""
    return -0.5 + ALPHA * u_exact(x)


def g_west(x):
    """Robin data on the west side, where ``u = 1``."""
    return 0.5 + ALPHA * u_exact(x)


def g_horizontal(x):
    """Robin data on south and north, where ``grad u . n = 0``."""
    return ALPHA * u_exact(x)


class MixedRobin(AbstractPhysicalWeakModel):
    """Dirichlet west, Robin east, homogeneous Neumann south/north."""

    def __init__(self):
        super().__init__(dim=DIM)
        self.weak_forms = {
            "interior": LaplacianWeakForm(dim=DIM, f=lambda x: jnp.zeros(OUT_DIM))
        }
        self.boundary_conditions = {
            "west": Dirichlet(u_exact),
            "east": Robin(ALPHA, g_east),
            "south": Neumann(lambda x: jnp.zeros(OUT_DIM)),
            "north": Neumann(lambda x: jnp.zeros(OUT_DIM)),
        }


class PureRobin(AbstractPhysicalWeakModel):
    """Robin on all four sides: no constrained DOF anywhere."""

    def __init__(self):
        super().__init__(dim=DIM)
        self.weak_forms = {
            "interior": LaplacianWeakForm(dim=DIM, f=lambda x: jnp.zeros(OUT_DIM))
        }
        self.boundary_conditions = {
            "west": Robin(ALPHA, g_west),
            "east": Robin(ALPHA, g_east),
            "south": Robin(ALPHA, g_horizontal),
            "north": Robin(ALPHA, g_horizontal),
        }


# %% Mesh, basis, scheme.
def make_scheme(model):
    """A fresh scheme for one model, so each run starts from zero DOFs."""
    mesh = Mesh(
        dim=DIM,
        n_cells=(N_CELLS, N_CELLS),
        ref_quad=UnitSquareTensorized(dim=DIM, order=QUAD_ORDER),
        mapping=Mapping(mappings=[InvertibleFunction(lambda x: x, lambda y: y)]),
    )
    basis = AnalyticBasis(
        nb_basis=(POLY_ORDER + 1) ** DIM,
        out_dim=OUT_DIM,
        mesh=mesh,
        local_basis=lambda c, i, m: local_lagrange_basis(
            c, i, m, order=POLY_ORDER, out_dim=OUT_DIM
        ),
        basis_type="scalar",
    )
    return EllipticFEscheme(model, VariablesFE(basis=basis, nb_variables=OUT_DIM))


print(f"FEM Q{POLY_ORDER} on {N_CELLS}x{N_CELLS}, alpha = {ALPHA}")
for name, model in (
    ("mixed Dirichlet/Robin/Neumann", MixedRobin),
    ("pure Robin", PureRobin),
):
    for tag, matrix_free in (("assembled J", False), ("matrix-free CG", True)):
        scheme = EllipticFEscheme.solve(
            make_scheme(model()), matrix_free=matrix_free, max_iter=1
        )
        err = float(l2_error(scheme, u_exact, relative=True))
        print(f"  [{name:29s} | {tag:14s}] relative L2 error = {err:.3e}")

print(
    "the exact solution is affine, so these should sit at solver tolerance, "
    "not merely converge"
)

# %% The pure-Robin solution: no constrained DOF anywhere, the boundary
# integral alone pins it down.
scheme = EllipticFEscheme.solve(make_scheme(PureRobin()), max_iter=1)
plot_solution_2d(
    scheme,
    u_exact,
    title=f"FEM Q{POLY_ORDER}, {N_CELLS}x{N_CELLS}: Robin on all four sides "
    f"(alpha={ALPHA}), no constrained DOF",
)
# grad u . n = 0 on south and north, so there too the slices collapse onto
# 1 - x/2 -- here held in place by the Robin term alone.
plot_boundary_profiles(
    scheme,
    u_exact,
    title="Pure Robin: every y-slice collapses onto $1-x/2$",
)