"""DG on the unit square: Laplace with Dirichlet, Neumann and Robin conditions.

    -Delta u = 0   in (0,1)^2,   one condition per side, keyed by side name

Usage: ``python solve_2d_laplacian_boundary_conditions.py [--bc CASE]``, with
``CASE`` one of

* ``dirichlet_neumann`` (default): ``u = 1`` west, ``u = 0`` east,
  ``du/dn = 0`` south and north; exact solution ``u = 1 - x``;
* ``dirichlet_robin``: ``u`` given west, ``du/dn + alpha u = g`` east,
  ``du/dn = 0`` south and north; exact solution ``u = 1 - x/2``;
* ``robin``: ``du/dn + alpha u = g`` on all four sides, ``alpha = 2``; same
  exact solution, and no constrained DOF anywhere: the boundary integral alone
  pins the solution down.

The exact solutions are affine, so they lie in every polynomial space: the
error measures the SCHEME, not the approximation, and must sit at solver
tolerance -- anything above is a bug in the boundary handling. The Robin
solution is ``1 - x/2`` rather than ``1 - x`` so that it does not vanish on
the Robin sides: with ``u = 0`` there the ``alpha u`` term would be zero and
the case would only check Neumann.

In DG a prescribed flux (Neumann, Robin) enters the residual as ``-(flux) v``
whatever the numerical flux; only Dirichlet goes through the numerical flux
(DG has no boundary DOF, it uses a ghost state), which is why the pure-Robin
case never calls the flux's boundary branch. The conditions are keyed by side
name, as a PINN model keys its residuals by boundary label
(``Mesh.boundary_groups``).
"""

# %%
import argparse

import jax.numpy as jnp

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 l2_error
from scimba_jax.linear_approximation.galerkin.dg.elliptic_dg_scheme import (
    EllipticDGscheme,
)
from scimba_jax.linear_approximation.galerkin.dg.flux import SIPGFlux
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.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
ALPHA = 2.0


# ── Data: module-level functions, a stable compile key ────────────────────────


def identity(x):
    return x


def taylor(coords, i, mesh):
    return local_taylor_basis(coords, i, mesh, order=POLY_ORDER, out_dim=1)


def zero(x):
    return jnp.zeros(1)


def one(x):
    return jnp.ones(1)


def u_dirichlet_neumann(x):
    """``u = 1 - x``."""
    return jnp.array([1.0 - x[0]])


def u_robin(x):
    """``u = 1 - x/2``, nonzero on every Robin side."""
    return jnp.array([1.0 - 0.5 * x[0]])


# grad u = (-1/2, 0) for u_robin: grad u . n is -1/2 east, +1/2 west, 0 on the
# horizontal sides, and the Robin datum is g = grad u . n + alpha u there.
def g_east(x):
    return -0.5 + ALPHA * u_robin(x)


def g_west(x):
    return 0.5 + ALPHA * u_robin(x)


def g_horizontal(x):
    return ALPHA * u_robin(x)


CASES = {
    "dirichlet_neumann": (
        u_dirichlet_neumann,
        {"west": Dirichlet(one), "east": Dirichlet(zero)},
    ),
    "dirichlet_robin": (
        u_robin,
        {"west": Dirichlet(u_robin), "east": Robin(ALPHA, g_east)},
    ),
    "robin": (
        u_robin,
        {
            "west": Robin(ALPHA, g_west),
            "east": Robin(ALPHA, g_east),
            "south": Robin(ALPHA, g_horizontal),
            "north": Robin(ALPHA, g_horizontal),
        },
    ),
}


def model(case):
    """Laplace, the case's conditions, homogeneous Neumann on the other sides."""
    conditions = {"south": Neumann(zero), "north": Neumann(zero)}
    conditions.update(CASES[case][1])
    physical = AbstractPhysicalWeakModel(dim=DIM)
    physical.add_weak_form("interior", LaplacianWeakForm(dim=DIM, f=zero))
    for side, condition in conditions.items():
        physical.add_boundary_condition(side, condition)
    return physical


def scheme(case):
    mesh = Mesh(
        dim=DIM,
        n_cells=(N_CELLS, N_CELLS),
        ref_quad=UnitSquareTensorized(dim=DIM, order=QUAD_ORDER),
        mapping=Mapping(mappings=[InvertibleFunction(identity, identity)]),
    )
    basis = AnalyticBasis(
        nb_basis=(POLY_ORDER + 1) ** DIM,
        out_dim=1,
        mesh=mesh,
        local_basis=taylor,
        basis_type="scalar",
    )
    return EllipticDGscheme(
        model(case),
        VariablesDG(basis=basis, nb_variables=1),
        SIPGFlux(sigma=POLY_ORDER * (POLY_ORDER + 1) * DIM, h=1.0 / N_CELLS),
    )


# %%
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    parser.add_argument("--bc", choices=sorted(CASES), default="dirichlet_neumann")
    case = parser.parse_args().bc
    u_exact = CASES[case][0]

    solved = EllipticDGscheme.solve(scheme(case), max_iter=1)
    error = float(l2_error(solved, u_exact, relative=True))
    print(f"DG Q{POLY_ORDER} on {N_CELLS}x{N_CELLS}, --bc {case}")
    print(f"relative L2 error = {error:.3e}  (affine solution: solver tolerance)")

    plot_solution_2d(solved, u_exact, title=f"DG Q{POLY_ORDER}, {case}")
    # du/dn = 0 on south and north: every y-slice collapses onto u_exact(x).
    plot_boundary_profiles(
        solved, u_exact, title=f"{case}: every y-slice collapses onto u(x)"
    )