"""The far-field conditions on the two classical exterior problems, CG-FEM Q2.

A far-field condition replaces the unbounded exterior of a domain by its exact
Dirichlet-to-Neumann map on an artificial boundary. The classical check is an
exterior problem with a closed-form solution: the artificial boundary must then
be invisible -- the discrete solution as good as if the exact trace had been
imposed there.

* **Laplace, outside a disk.** ``-Laplace(u) = 0`` for ``r > r0``, ``u =
  cos(2 t)`` on ``r = r0``, ``u`` bounded at infinity: ``u = (r0/r)^2 cos(2 t)``.
  Solved on the annulus ``r0 < r < R`` with :class:`LaplaceCircleFarField` on
  ``r = R``.
* **Grad-Shafranov, outside a sphere.** ``-div((1/r) grad psi) = 0`` in the
  poloidal half plane outside ``rho = rho0``, ``psi = 0`` on the axis, the
  dipole's flux ``psi = r^2 / rho^3`` on ``rho = rho0``, ``psi -> 0`` at
  infinity: ``psi = r^2 / rho^3`` everywhere. Solved on the half annulus
  ``rho0 < rho < R`` with :class:`GradShafranovFarField` on the outer arc.

Neither involves a source or a Green function: the exact solution is a formula.
Three conditions on the outer boundary are compared on three meshes: the far
field, the exact trace (the best any condition can do), and ``u = 0`` (a
boundary pushed "far enough"). Needs the optional ``mesh`` extra (gmsh).
"""

# %%
from pathlib import Path

import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.tri import Triangulation

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.variables.variables_fe import VariablesFE
from scimba_jax.mapping.macro_mesh import (
    ExactCircle,
    ExactHalfAnnulus,
    macro_mesh_from_loops,
)
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,
)
from scimba_jax.physical_models.classical_weakform.laplacian_weak_form import (
    LaplacianWeakForm,
)
from scimba_jax.physical_models.elliptic_pde.grad_shafranov_free_boundary import (
    GradShafranovFarField,
)
from scimba_jax.physical_models.weak_boundary_conditions import (
    Dirichlet,
    LaplaceCircleFarField,
)

ORDER = 2
R_IN, R_OUT = 0.5, 1.5
MESH_SIZES = (0.2, 0.1, 0.05)


# ── Exact solutions ─────────────────────────────────────────────────────────


def laplace_exact(x):
    """``(r0/r)^2 cos(2 t) = r0^2 (x^2 - y^2) / r^4``."""
    r2 = x[..., 0] ** 2 + x[..., 1] ** 2
    return R_IN**2 * (x[..., 0] ** 2 - x[..., 1] ** 2) / r2**2


def dipole_exact(x):
    """``r^2 / rho^3``, the flux of a magnetic dipole on the axis."""
    return x[..., 0] ** 2 / (x[..., 0] ** 2 + x[..., 1] ** 2) ** 1.5


def laplace_trace(x):
    return jnp.atleast_1d(laplace_exact(x))


def dipole_trace(x):
    return jnp.atleast_1d(dipole_exact(x))


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


def inverse_radius_identity(x):
    return jnp.eye(2) / x[0]


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


def no_reaction(x):
    return jnp.zeros(())


def lagrange(y, i, mesh):
    return local_lagrange_basis(y, i, mesh, order=ORDER, out_dim=1)


def lagrange_by_logical(y, i, mesh):
    return local_lagrange_basis_by_logical(y, i, mesh, order=ORDER, out_dim=1)


# ── The two cases ───────────────────────────────────────────────────────────


def laplace_case(mesh_size):
    """Annulus: inner circle, outer circle."""
    macro = macro_mesh_from_loops(
        ExactCircle(0.0, 0.0, R_OUT),
        holes=[ExactCircle(0.0, 0.0, R_IN)],
        order=ORDER,
        mesh_size=mesh_size,
    )
    mesh = UnstructuredMesh.from_macro_mesh(macro, UnitSquareTensorized(dim=2, order=4))
    middle = 0.5 * (R_IN + R_OUT)
    mesh.label_boundary("inner", lambda x: np.linalg.norm(x, axis=1) < middle)
    mesh.label_boundary("outer", lambda x: np.linalg.norm(x, axis=1) > middle)
    weak_form = LaplacianWeakForm(dim=2, f=zero)
    conditions = {"inner": Dirichlet(laplace_trace)}
    outer = {
        "far field": LaplaceCircleFarField(),
        "exact trace": Dirichlet(laplace_trace),
        "u = 0": Dirichlet(zero),
    }
    return mesh, weak_form, conditions, outer, laplace_exact


def dipole_case(mesh_size):
    """Half annulus: the axis, the inner half circle, the outer arc."""
    macro = macro_mesh_from_loops(
        ExactHalfAnnulus(R_IN, R_OUT), order=ORDER, mesh_size=mesh_size
    )
    mesh = UnstructuredMesh.from_macro_mesh(macro, UnitSquareTensorized(dim=2, order=4))
    middle = 0.5 * (R_IN + R_OUT)
    mesh.label_boundary("axis", lambda x: x[:, 0] < 1e-8)
    mesh.label_boundary(
        "inner", lambda x: (x[:, 0] >= 1e-8) & (np.linalg.norm(x, axis=1) < middle)
    )
    mesh.label_boundary(
        "outer", lambda x: (x[:, 0] >= 1e-8) & (np.linalg.norm(x, axis=1) > middle)
    )
    weak_form = EllipticWeakForm(
        dim=2, A=inverse_radius_identity, b=no_advection, c=no_reaction, f=zero
    )
    conditions = {"axis": Dirichlet(zero), "inner": Dirichlet(dipole_trace)}
    outer = {
        "far field": GradShafranovFarField(R_OUT),
        "exact trace": Dirichlet(dipole_trace),
        "u = 0": Dirichlet(zero),
    }
    return mesh, weak_form, conditions, outer, dipole_exact


def solve(mesh, weak_form, conditions, outer_condition):
    """Solve on ``mesh``; returns the node positions and the nodal solution."""
    model = AbstractPhysicalWeakModel(dim=2)
    model.add_weak_form("main", weak_form)
    for name, condition in conditions.items():
        model.add_boundary_condition(name, condition)
    model.add_boundary_condition("outer", outer_condition)
    basis = AnalyticBasis(
        nb_basis=(ORDER + 1) ** 2,
        out_dim=1,
        mesh=mesh,
        local_basis=lagrange,
        local_basis_by_logical=lagrange_by_logical,
        basis_type="scalar",
    )
    variables = VariablesFE(
        basis=basis, nb_variables=1, dof_map=UnstructuredLagrangeDofMap
    )
    scheme = EllipticFEscheme.solve(EllipticFEscheme(model, variables))
    nodes = np.asarray(
        variables.boundary_dof_positions(np.arange(variables.dofsl.shape[0]))
    )
    return nodes, np.asarray(scheme.variables.dofsl[:, 0])


# %% Convergence: max nodal error, relative to max |u|.
finest = {}
for title, case in (
    ("Laplace outside a disk, u = (r0/r)^2 cos 2t", laplace_case),
    ("Grad-Shafranov outside a sphere, psi = r^2 / rho^3", dipole_case),
):
    print(f"\n{title}")
    errors = {}
    for mesh_size in MESH_SIZES:
        mesh, weak_form, conditions, outer, exact = case(mesh_size)
        for label, condition in outer.items():
            nodes, u = solve(mesh, weak_form, conditions, condition)
            reference = np.asarray(exact(nodes))
            error = np.max(np.abs(u - reference)) / np.max(np.abs(reference))
            errors.setdefault(label, []).append(error)
            if label == "far field" and mesh_size == MESH_SIZES[-1]:
                finest[title] = (nodes, u, reference)
    print(f"{'outer condition':12s} " + " ".join(f"h={h:<8}" for h in MESH_SIZES))
    for label, values in errors.items():
        print(f"{label:12s} " + " ".join(f"{e:.2e}  " for e in values))

# %% Plot: the far-field solution on the finest mesh, and its error.
fig, axes = plt.subplots(2, 2, figsize=(11, 10))
for row, (title, (nodes, u, reference)) in enumerate(finest.items()):
    # Delaunay on the nodes fills the hole around the origin: mask it.
    triangulation = Triangulation(nodes[:, 0], nodes[:, 1])
    centres = nodes[triangulation.triangles].mean(axis=1)
    triangulation.set_mask(np.linalg.norm(centres, axis=1) < R_IN)
    for ax, field, name in (
        (axes[row, 0], u, "far-field solution"),
        (axes[row, 1], np.abs(u - reference), "|error|"),
    ):
        tc = ax.tricontourf(triangulation, field, levels=40, cmap="viridis")
        fig.colorbar(tc, ax=ax)
        ax.set_aspect("equal")
        ax.set_title(f"{name}\n{title}", fontsize=9)
fig.tight_layout()
out = Path(__file__).with_name("exterior_problems.png")
fig.savefig(out, dpi=120)
print(f"\nSaved figure to {out}")
plt.show()