"""Complex Helmholtz on an unstructured disk, with a point source and an ABC.

    Delta u + kappa u = f          in the unit disk, u = (u_Re, u_Im)
    d_n u - i*k*u = 0              on the whole boundary

``kappa = k^2`` and ``f`` is a regularized Dirac (a narrow Gaussian bump) at a
point away from the center, so the solution looks like an outgoing wave
radiating from the source and leaving through the boundary -- which is exactly
what the first-order Sommerfeld/absorbing condition (ABC) ``d_n u - i k u = 0``
approximates. This exercises three things at once that the other FEM examples
keep separate: a genuinely complex-valued field, a Robin condition with a
complex coefficient, and a *traditional* unstructured mesh (freely triangulated
then recombined into quads by GMSH, not the block-structured O-grid used in
``solve_laplacian_unstructured_2d.py``/``solve_hdiv_hcurl_unstructured_disk_2d.py``).

**Complex = a real pair.** Nothing in the FEM/DG stack has a complex dtype;
the library-wide convention (already used by ``HelmholtzWeakForm``, the PINN
``HelmholtzND``, and the DG ``HelmholtzSIPGRobinFlux``) is ``u = (u_Re, u_Im)``
as an ``out_dim=2`` ``"vec"`` space, with the interior weak form ``eq_re``/
``eq_im`` decoupled (kappa is real) and only the boundary term coupling them
(a complex product mixes real and imaginary parts).

**The one missing piece.** ``physical_models.weak_boundary_conditions.Robin``
has no complex coefficient, and -- more subtly -- it is written for the
``EllipticWeakForm``/``LaplacianWeakForm`` sign convention (``-Delta u = f``,
bilinear form ``+grad u . grad v``). ``HelmholtzWeakForm`` uses the opposite
one: integration by parts of ``Delta u + kappa u = f`` gives bilinear form
``-grad u . grad v + kappa u v``, which flips the sign the natural boundary
term must carry. ``ComplexRobin`` below re-derives it from scratch for that
convention (see its docstring) rather than reusing ``Robin`` with a wrong
sign. Everything else -- ``HelmholtzWeakForm`` itself, ``AnalyticBasis``,
``UnstructuredMesh``, ``UnstructuredLagrangeDofMap``, ``EllipticFEscheme`` --
is reused unchanged.

Needs the optional ``mesh`` extra (pygmsh/meshio/gmsh). Saves
``helmholtz_robin_unstructured_disk.png`` next to this file and shows it.
"""

# %%
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,
    local_lagrange_basis_by_logical,
)
from scimba_jax.linear_approximation.basis.dof_map import UnstructuredLagrangeDofMap
from scimba_jax.linear_approximation.basis.general_bases import AnalyticBasis
from scimba_jax.linear_approximation.galerkin.fem.elliptic_fe_scheme import (
    EllipticFEscheme,
)
from scimba_jax.linear_approximation.meshes.unstructured_mesh import UnstructuredMesh
from scimba_jax.linear_approximation.quad.gauss_quad import UnitSquareTensorized
from scimba_jax.linear_approximation.solvers import LinearSolve
from scimba_jax.linear_approximation.variables.variables_fe import VariablesFE
from scimba_jax.mapping.macro_mesh import macro_mesh_from_curve
from scimba_jax.physical_models.abstract_physical_weak_model import (
    AbstractPhysicalWeakModel,
)
from scimba_jax.physical_models.classical_weakform.helmholtz_weak_form import (
    HelmholtzWeakForm,
)
from scimba_jax.physical_models.weak_boundary_conditions import ComplexRobin

_HERE = Path(__file__).resolve().parent

DIM = 2
ORDER = 3  # geometric AND FE degree: isoparametric, like the O-grid disk example
MESH_SIZE = 0.1
QUAD_ORDER = 2 * ORDER + 2

K = 32.0  # wavenumber
KAPPA = K**2

# source location, off-center for a visible pattern
X0 = jnp.array([[0.3, -0.2], [-0.4, 0.5]])
# regularization width of the Dirac, resolved by MESH_SIZE/QUAD_ORDER
SIGMA_SOURCE = 0.01


# %% The regularized Dirac: a narrow, L1-normalized real Gaussian bump.
def f_source(x):
    """A real point source, ``[f_Re, f_Im] = [delta_sigma(x - x0), 0]``."""
    bump = jnp.sum(jnp.exp(-jnp.sum((x - X0) ** 2, axis=1) / (2.0 * SIGMA_SOURCE**2)))
    bump = bump / (2.0 * jnp.pi * SIGMA_SOURCE**2)
    return jnp.array([bump, 0.0])


# %% The absorbing condition is the library's ComplexRobin (weak_boundary_conditions),
# the complex Robin of HelmholtzWeakForm's sign convention: bilinear -beta u v, linear
# -g v, both signs opposite to the real Robin (which pairs with -Delta u = f).
def g_zero(x):
    """Homogeneous ABC data."""
    return jnp.zeros(2)


# %% Mesh: a traditional unstructured mesh of the disk
def disk_boundary(t):
    return (np.cos(2.0 * np.pi * t), np.sin(2.0 * np.pi * t))


macro = macro_mesh_from_curve(disk_boundary, order=ORDER, mesh_size=MESH_SIZE)
mesh = UnstructuredMesh.from_macro_mesh(
    macro, UnitSquareTensorized(dim=DIM, order=QUAD_ORDER)
)
print(f"unstructured disk mesh: {mesh.n_cells_total} cells, degree {mesh.order}")

# %% out_dim=2 vector space for (u_Re, u_Im), same isoparametric setup as the
# O-grid disk examples.
basis = AnalyticBasis(
    nb_basis=(ORDER + 1) ** DIM,
    out_dim=2,
    mesh=mesh,
    basis_type="vec",
    local_basis=lambda y, i, m: local_lagrange_basis(y, i, m, order=ORDER, out_dim=2),
    local_basis_by_logical=lambda y, i, m: local_lagrange_basis_by_logical(
        y, i, m, order=ORDER, out_dim=2
    ),
)
variables = VariablesFE(basis=basis, nb_variables=2, dof_map=UnstructuredLagrangeDofMap)

model = AbstractPhysicalWeakModel(dim=DIM)
model.add_weak_form("main", HelmholtzWeakForm(dim=DIM, kappa=KAPPA, f=f_source))
model.add_boundary_condition(
    "boundary", ComplexRobin(beta_re=0.0, beta_im=-K, g=g_zero)
)

scheme = EllipticFEscheme(model, variables)
# The problem is linear (HelmholtzWeakForm and ComplexRobin are both linear in
# u): one assembled solve, named as such rather than as max_iter=1. Also
# sidesteps matrix-free Newton-Krylov's default CG, which assumes a definite
# operator -- complex Helmholtz (kappa > 0) is not one.
scheme = EllipticFEscheme.solve(scheme, solver=LinearSolve(), verbose=True)
print(f"k = {K}, kappa = k^2 = {KAPPA}, {variables.n_nodes_total} nodes")


# %% Plot Re(u), Im(u), |u|. ``sample_solution`` is cell-by-cell with no point
# search, which is what a non-convex/curved domain needs (a rectangular grid
# would put points outside the disk) -- but it hardcodes component 0, so a
# complex field needs this small variant that keeps both components.
def sample_complex_solution(scheme, n_side=10):
    mesh = scheme.variables.mesh
    grid = jnp.stack(
        jnp.meshgrid(
            jnp.linspace(0.0, 1.0, n_side),
            jnp.linspace(0.0, 1.0, n_side),
            indexing="ij",
        ),
        axis=-1,
    ).reshape(-1, 2)
    dofs = scheme.variables.dofsl
    connectivity = scheme.variables.connectivity

    def one_cell(cell):
        points = mesh._unit_hypercube_to_cell(cell, grid)
        theta = dofs[connectivity[cell]]
        values = jax.vmap(
            lambda p: jnp.einsum(
                "iv,iv->v", theta, scheme.variables.trial_basis(cell, p[None, :])[0]
            )
        )(points)
        return points, values

    points, values = jax.vmap(one_cell)(jnp.arange(mesh.n_cells_total))

    corner = np.arange(n_side - 1)
    i, j = np.meshgrid(corner, corner, indexing="ij")
    bottom_left = (i * n_side + j).ravel()
    local = np.concatenate(
        [
            np.stack([bottom_left, bottom_left + n_side, bottom_left + 1], axis=1),
            np.stack(
                [bottom_left + n_side, bottom_left + n_side + 1, bottom_left + 1],
                axis=1,
            ),
        ]
    )
    offsets = (np.arange(mesh.n_cells_total) * n_side**2)[:, None, None]
    triangles = (local[None, :, :] + offsets).reshape(-1, 3)

    return (
        np.asarray(points).reshape(-1, 2),
        triangles,
        np.asarray(values).reshape(-1, 2),
    )


points, triangles, values = sample_complex_solution(scheme, n_side=10)
u_re, u_im, amplitude = values[:, 0], values[:, 1], np.hypot(values[:, 0], values[:, 1])

figure, axes = plt.subplots(1, 3, figsize=(15, 4.6))
for ax, field, label in zip(axes, (u_re, u_im, amplitude), ("Re(u)", "Im(u)", "|u|")):
    drawing = ax.tricontourf(
        points[:, 0], points[:, 1], triangles, field, levels=40, cmap="turbo"
    )
    figure.colorbar(drawing, ax=ax, fraction=0.046)
    ax.set_title(label)
    ax.set_aspect("equal")
    ax.set_xlabel("x")
    ax.set_ylabel("y")

figure.suptitle(
    f"Complex Helmholtz, k={K}, unstructured disk ({mesh.n_cells_total} cells): "
    "regularized Dirac source + absorbing (complex Robin) boundary"
)
figure.tight_layout()
figure.savefig(_HERE / "helmholtz_robin_unstructured_disk.png", dpi=130)
print(f"\nSaved {_HERE / 'helmholtz_robin_unstructured_disk.png'}")
plt.show()