"""CG-FEM on the unit square with mixed boundary conditions, one per side.

    -Delta u = 0    in (0,1)^2
     u = 1          on the west side       (Dirichlet, strongly imposed)
     u = 0          on the east side       (Dirichlet, strongly imposed)
     du/dn = 0      on south and north     (Neumann, natural)

whose solution is exactly ``u(x, y) = 1 - x``. Being affine, it lies in every
polynomial space used here, so the error measures the *scheme*, not the
approximation: anything above solver tolerance is a bug in the boundary
handling, which makes this a far sharper check than a convergence rate.

This is the CG-FEM twin of the DG example of the same name, with the *same*
model class: the conditions are keyed by side name, exactly as a PINN model
keys its residuals by boundary label, and the scheme decides how to enforce
them. Here that means constraining the west/east nodes and integrating the
prescribed flux on the south/north faces — a policy of the scheme, invisible
in the physics.
"""

# %%
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
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


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


class MixedLaplacian(AbstractPhysicalWeakModel):
    """Laplace with Dirichlet east/west and 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(lambda x: jnp.ones(OUT_DIM)),
            "east": Dirichlet(lambda x: jnp.zeros(OUT_DIM)),
            "south": Neumann(lambda x: jnp.zeros(OUT_DIM)),
            "north": Neumann(lambda x: jnp.zeros(OUT_DIM)),
        }


# %% Mesh, basis, scheme.
def make_scheme():
    """A fresh scheme, so each solver path 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(
        MixedLaplacian(), VariablesFE(basis=basis, nb_variables=OUT_DIM)
    )


mesh_sides = sorted(make_scheme().variables.mesh.boundary_groups)
print(f"FEM Q{POLY_ORDER} on {N_CELLS}x{N_CELLS}, sides: {mesh_sides}")

# %% Both solver paths must see the same boundary terms: the assembled Jacobian
# and the matrix-free JVP are written separately, so agreeing here is itself a
# check that the face loop was wired into both.
#
# The matrix-free path is a real test here, and it was not before: it runs CG,
# which needs a symmetric operator, and this is the first FEM example with
# inhomogeneous Dirichlet data (u = 1 on the west side). The scheme imposes
# Dirichlet by substituting the data into the constrained DOFs rather than by
# overwriting rows, which keeps the Jacobian symmetric -- see _lift_dirichlet.
for tag, matrix_free in (("assembled J", False), ("matrix-free CG", True)):
    scheme = EllipticFEscheme.solve(make_scheme(), matrix_free=matrix_free, max_iter=1)
    err = float(l2_error(scheme, u_exact, relative=True))
    print(f"  [{tag}] relative L2 error = {err:.3e}")

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

# %% What the conditions did, rather than just how small the error is.
plot_solution_2d(
    scheme,
    u_exact,
    title=f"CG-FEM Q{POLY_ORDER}, {N_CELLS}x{N_CELLS}: u=1 west, u=0 east, "
    "du/dn=0 south/north",
)
# Homogeneous Neumann on south and north means u does not vary with y, so all
# slices must collapse onto the single line 1 - x. They would visibly fan out
# if the natural condition were mishandled.
plot_boundary_profiles(
    scheme,
    u_exact,
    title="Neumann south/north: every y-slice collapses onto $1-x$",
)