"""Convection-dominated DG: a boundary layer, SIPG diffusion plus upwind advection.

Problem (1D, boundary layer at x = 1):

  d_x u - (1/Pe) d_xx u = r  on (0, 1), homogeneous Dirichlet,
  u_exact(x) = r * (x - (exp(Pe x) - 1) / (exp(Pe) - 1)),
  (r, Pe) = (1.5, 90)  (test case mu(2) of the reference paper).

This is ``-div(A grad u) + b.grad u + c u = f`` with ``A = 1/Pe``, ``b = 1``,
``c = 0``, ``f = r``, discretised in a broken Q1 space. The diffusion faces get
the symmetric interior penalty (``SIPGFlux``), the advection faces the upwind
term (``UpwindFlux``), the two composed by ``SumFlux``.

scimba writes advection in NON-conservative form (a pure volume term), so the
upwind face term carries both a consistency coupling and a jump penalty; the
penalty alone is not upwinding and makes a boundary layer worse, not better.
The comparison of the face terms (none, centred, upwind), of the regimes
(diffusion, diffusion-advection(-reaction), Pe = 90 and Pe = 500), of the
bases and of the degrees lives in the benchopt benchmark
``benchmarks/benchmarks_jax/dg_convection_diffusion``.
"""

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

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.dg.elliptic_dg_scheme import (
    EllipticDGscheme,
)
from scimba_jax.linear_approximation.galerkin.dg.flux import (
    SIPGFlux,
    SumFlux,
    UpwindFlux,
)
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.diffusion_advection_reaction_weak_form import (  # noqa: E501
    EllipticWeakForm,
)

physical_dim = 1
out_dim = 1
poly_deg = 1
n_cells = 32

r_val = 1.5  # constant source term
pe_val = 90.0  # Peclet number (convection / diffusion)


# ── The problem: module-level callables, a stable compile key ────────────────


def identity(x):
    """The identity mapping of (0, 1)."""
    return x


def diffusion_tensor(x):
    """``A = (1/Pe) I``."""
    return (1.0 / pe_val) * jnp.eye(physical_dim)


def advection_field(x):
    """``b = 1``."""
    return jnp.array([1.0])


def reaction(x):
    """No reaction."""
    return jnp.array(0.0)


def source(x):
    """``f = r``."""
    return jnp.array(r_val)


def dirichlet_bc(x):
    """Homogeneous Dirichlet data."""
    return jnp.zeros(out_dim)


def local_basis(coords, i, mesh):
    """The Q1 Lagrange basis of a cell."""
    return local_lagrange_basis(coords, i, mesh, order=poly_deg, out_dim=out_dim)


def u_exact_fn(x):
    """Exact solution: linear ramp plus an exponential layer at x = 1."""
    return r_val * (x[0] - (jnp.exp(pe_val * x[0]) - 1.0) / (jnp.exp(pe_val) - 1.0))


pde = EllipticWeakForm(
    dim=physical_dim, A=diffusion_tensor, b=advection_field, c=reaction, f=source
)

mesh = Mesh(
    dim=physical_dim,
    n_cells=[n_cells],
    ref_quad=UnitSquareTensorized(dim=physical_dim, order=2 * poly_deg + 1),
    mapping=Mapping(mappings=[InvertibleFunction(identity, identity)]),
)
basis = AnalyticBasis(
    nb_basis=poly_deg + 1,
    out_dim=out_dim,
    mesh=mesh,
    local_basis=local_basis,
    basis_type="scalar",
)
scheme = EllipticDGscheme(
    AbstractPhysicalWeakModel.from_weak_form(pde, dirichlet=dirichlet_bc),
    VariablesDG(basis=basis, nb_variables=out_dim),
    SumFlux([SIPGFlux(sigma=poly_deg * (poly_deg + 1) * physical_dim), UpwindFlux()]),
)

# ⚠ matrix_free=False: the system is NON SYMMETRIC (advection), and the
# matrix-free Krylov defaults to CG ("lax", SPD only), which stalls or diverges
# on it. The assembled Jacobian is solved directly.
solved = EllipticDGscheme.solve(scheme, matrix_free=False)

err = float(l2_error(assembler=solved, u_exact_fn=u_exact_fn, relative=True))
print(
    f"Pe = {pe_val:.0f}, Q{poly_deg}, {n_cells} cells "
    f"(cell Pe = {pe_val / n_cells:.1f}): relative L2 error {err:.3e}"
)

# ── The solution, cell by cell (a DG solution is discontinuous) ───────────────
fig, (ax_sol, ax_err) = plt.subplots(1, 2, figsize=(11, 4.5))
x_fine = jnp.linspace(0.0, 1.0, 1000)[:, None]
ax_sol.plot(x_fine[:, 0], jax.vmap(u_exact_fn)(x_fine), "k--", lw=2, label="exact")

ref_pts = jnp.linspace(0.0, 1.0, 30)[:, None]
cell_mesh = solved.variables.mesh
x_cells = jax.vmap(
    lambda c: cell_mesh.mapping.local_mapping(
        cell_mesh._unit_hypercube_to_cell(c, ref_pts)
    )
)(cell_mesh.cells_idx)
u_cells = solved.variables.evaluate_quad(x_cells)
u_ex_cells = jax.vmap(jax.vmap(u_exact_fn))(x_cells).reshape(u_cells.shape)
for c in range(x_cells.shape[0]):
    label = "SIPG + upwind" if c == 0 else None
    ax_sol.plot(x_cells[c, :, 0], u_cells[c, :, 0], "tab:blue", lw=1.5, label=label)
    ax_err.semilogy(
        x_cells[c, :, 0],
        jnp.abs(u_cells[c, :, 0] - u_ex_cells[c, :, 0]),
        "tab:blue",
        lw=1.2,
    )
ax_sol.set_title(f"Solution, {n_cells} cells, Pe = {pe_val:.0f}")
ax_sol.set_xlabel("x")
ax_sol.legend()
ax_err.set_title(f"Pointwise error (relative L2 {err:.2e})")
ax_err.set_xlabel("x")
ax_err.grid(alpha=0.3)
fig.tight_layout()

plt.show()