"""Helmholtz with a perfectly matched layer, FEM and DG, against the exact field.

A Gaussian source radiates in free space; the square ``[0, 1]^2`` is
surrounded by a PML of thickness ``DELTA`` (Bérenger 1994) closed by a
Dirichlet condition:

    div(Λ grad u) + k^2 J u = J f      on [-δ, 1+δ]^2,   u = 0 on its boundary

with ``Λ = diag(s_2/s_1, s_1/s_2)``, ``J = s_1 s_2`` and ``s_j = 1 + i σ_j/k``
the complex stretching of each coordinate outside the square
(:class:`PerfectlyMatchedLayer`, calibrated to a reflection ``R0``). Inside the
square ``Λ = I``, ``J = 1``: the equation is the physical one.

**The exact solution.** For a radial source ``f(x) = exp(-|x - x0|^2 / 2σ^2)``
the free-space field is, outside the source's support, a pure outgoing
Hankel wave (Graf's addition theorem):

    u(x) = -(i/4) H0^(1)(k r) 2π σ^2 exp(-k^2 σ^2 / 2),   r = |x - x0|

so the error measured in the square, away from the source, is the whole error
-- the layer's reflections and the discretisation, nothing else.

**The solve.** BiCGSTAB preconditioned by one multigrid V-cycle on the
shifted operator ``div(Λ grad) + (k^2 + i ε) J``, ``ε = 2 k``
(:func:`shifted_laplacian_solver`): the same factory builds the problem at
``shift = 0`` and the preconditioner's operator at ``shift = ε``, for the FEM
and for the DG discretisation (symmetric interior penalty,
:class:`HelmholtzSIPGFlux`, which reads the layer's ``Λ``).

Needs no optional extra (structured mesh). Saves
``solve_2d_helmholtz_pml_fem_dg.png``.
"""

# %%
import time
from pathlib import Path

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

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.galerkin.dg.elliptic_dg_scheme import (
    EllipticDGscheme,
)
from scimba_jax.linear_approximation.galerkin.dg.flux import HelmholtzSIPGFlux
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.solvers.helmholtz import shifted_laplacian_solver
from scimba_jax.linear_approximation.variables.variables_dg import VariablesDG
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.helmholtz_weak_form import (
    HelmholtzWeakForm,
    PerfectlyMatchedLayer,
)
from scimba_jax.physical_models.weak_boundary_conditions import Dirichlet
from scimba_jax.utils.bessel import hankel1_0

jax.config.update("jax_enable_x64", True)

K = 30.0  # wavenumber: 4.8 wavelengths across the square
SIGMA = 0.03  # width of the Gaussian source
X0 = jnp.array([0.5, 0.5])
DELTA, R0 = 0.3, 1e-6  # layer thickness and target reflection
ORDER = 2  # Q2, FEM and DG
MESHES = (24, 32, 40)  # cells per side of [-δ, 1+δ]^2 (the MG's coarse level is half)
EXCLUDE = 0.15  # the error is measured at r > EXCLUDE from the source
PML = PerfectlyMatchedLayer((0.0, 1.0, 0.0, 1.0), DELTA, K, reflection=R0)
SIDES = ("west", "east", "south", "north")
LENGTH = 1.0 + 2.0 * DELTA


def source(x):
    return jnp.array([jnp.exp(-jnp.sum((x - X0) ** 2) / (2.0 * SIGMA**2)), 0.0])


def exact(x):
    """``-(i/4) H0(k r) C``: ``Re = C Y0 / 4``, ``Im = -C J0 / 4``."""
    weight = 2.0 * jnp.pi * SIGMA**2 * jnp.exp(-(K**2) * SIGMA**2 / 2.0)
    h = hankel1_0(K * jnp.linalg.norm(x - X0))
    return weight / 4.0 * jnp.array([h.imag, -h.real])


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


def to_domain(y):
    return -DELTA + LENGTH * y


def from_domain(x):
    return (x + DELTA) / LENGTH


def lagrange(c, i, mesh):
    return local_lagrange_basis(c, i, mesh, order=ORDER, out_dim=2)


def make_variables(kind, n_cells):
    mesh = Mesh(
        dim=2,
        n_cells=tuple(n_cells),
        ref_quad=UnitSquareTensorized(dim=2, order=ORDER + 3),
        mapping=Mapping(mappings=[InvertibleFunction(to_domain, from_domain)]),
        is_identity_mapping=False,
    )
    basis = AnalyticBasis(
        nb_basis=(ORDER + 1) ** 2,
        out_dim=2,
        mesh=mesh,
        basis_type="vec",
        local_basis=lagrange,
    )
    if kind == "FEM":
        return VariablesFE(basis=basis, nb_variables=2)
    return VariablesDG(basis=basis, nb_variables=2)


def fem_factory(n_cells, shift=0.0):
    return _factory("FEM", n_cells, shift)


def dg_factory(n_cells, shift=0.0):
    return _factory("DG", n_cells, shift)


def _factory(kind, n_cells, shift):
    """The problem at ``shift = 0``, the preconditioner's operator at ``shift = ε``."""
    model = AbstractPhysicalWeakModel(dim=2)
    model.add_weak_form(
        "main", HelmholtzWeakForm(2, K**2, source, kappa_im=shift, pml=PML)
    )
    for side in SIDES:
        model.add_boundary_condition(side, Dirichlet(zero))
    variables = make_variables(kind, n_cells)
    if kind == "FEM":
        return EllipticFEscheme(model, variables)
    flux = HelmholtzSIPGFlux(10.0 * ORDER * (ORDER + 1), h=LENGTH / n_cells[0])
    return EllipticDGscheme(model, variables, flux)


# %% Solve, and measure against the exact field.
rng = np.random.default_rng(0)
points = rng.uniform(0.0, 1.0, (20_000, 2))
points = jnp.asarray(points[np.linalg.norm(points - np.asarray(X0), axis=1) > EXCLUDE])
reference = np.asarray(jax.vmap(exact)(points))
factories = {
    "FEM": (EllipticFEscheme, fem_factory),
    "DG": (EllipticDGscheme, dg_factory),
}
results = {}
print(f"k = {K}, PML thickness {DELTA}, Q{ORDER}; error at r > {EXCLUDE} in [0,1]^2")
for kind, (cls, factory) in factories.items():
    for n in MESHES:
        start = time.perf_counter()
        solver = shifted_laplacian_solver(factory, (n, n), shift=2.0 * K, tol=1e-9)
        setup = time.perf_counter() - start
        start = time.perf_counter()
        solved, report = cls.solve(factory((n, n)), solver=solver, return_report=True)
        values = np.asarray(solved.variables.evaluate(points))
        solve_time = time.perf_counter() - start
        error = np.linalg.norm(values - reference) / np.linalg.norm(reference)
        results[(kind, n)] = (solved, error)
        n_dofs = solved.variables.dofsl.size
        print(
            f"  {kind:3s} {n:3d} x {n:<3d} ({n_dofs:6d} unknowns): rel. L2 error "
            f"{error:.2e}, {int(report.n_linear):3d} BiCGSTAB iterations"
            f"{'' if bool(report.converged) else ' (NOT CONVERGED)'}, "
            f"setup {setup:.1f} s, solve {solve_time:.1f} s"
        )
    errors = [results[(kind, n)][1] for n in MESHES]
    rates = np.log(np.array(errors[:-1]) / errors[1:]) / np.log(
        np.array(MESHES[1:]) / MESHES[:-1]
    )
    print(f"      orders {np.round(rates, 2)} (Q{ORDER}: {ORDER + 1} expected)")

# %% Figure: the field on the whole domain, the exact one, the errors.
side = np.linspace(-DELTA, 1.0 + DELTA, 241)
X, Y = np.meshgrid(side, side, indexing="ij")
grid = jnp.asarray(np.stack([X.ravel(), Y.ravel()], -1))
inside = (np.abs(X - 0.5) <= 0.5) & (np.abs(Y - 0.5) <= 0.5)
near = np.hypot(X - 0.5, Y - 0.5) < EXCLUDE
fine = MESHES[-1]
fields = {
    kind: np.asarray(results[(kind, fine)][0].variables.evaluate(grid))[:, 0].reshape(
        X.shape
    )
    for kind in factories
}
exact_grid = np.array(jax.vmap(exact)(grid))[:, 0].reshape(X.shape)
exact_grid[~inside | near] = np.nan
fig, axes = plt.subplots(1, 4, figsize=(21, 5), constrained_layout=True)
scale = np.nanmax(np.abs(exact_grid))
image = axes[0].pcolormesh(
    X, Y, fields["DG"], cmap="RdBu_r", vmin=-scale, vmax=scale, shading="auto"
)
axes[0].set_title(f"DG, Re u on the whole domain ({fine} x {fine})")
axes[1].pcolormesh(
    X, Y, exact_grid, cmap="RdBu_r", vmin=-scale, vmax=scale, shading="auto"
)
axes[1].set_title("exact, Re u (r > {:.2f})".format(EXCLUDE))
fig.colorbar(image, ax=axes[:2], shrink=0.8)
for ax, kind in zip(axes[2:], ("FEM", "DG")):
    error = np.abs(fields[kind] - exact_grid)
    err_image = ax.pcolormesh(
        X, Y, error, cmap="magma", vmin=0.0, vmax=np.nanmax(error), shading="auto"
    )
    ax.set_title(f"{kind}: |Re (u - u_exact)|, rel. L2 {results[(kind, fine)][1]:.1e}")
    fig.colorbar(err_image, ax=ax, shrink=0.8)
for ax in axes:
    ax.plot([0, 1, 1, 0, 0], [0, 0, 1, 1, 0], "k--", lw=0.8)
    ax.set_aspect("equal")
    ax.set_xlabel("x")
    ax.set_ylabel("y")
fig.suptitle(
    f"Helmholtz k = {K}, Gaussian source, PML of thickness {DELTA} (dashed: the physical square)"
)
out = Path(__file__).with_name("solve_2d_helmholtz_pml_fem_dg.png")
fig.savefig(out, dpi=120)
print(f"Saved figure to {out}")
plt.show()