"""2D magnetostatics with a permanent magnet, CG-FEM, materials as named regions.

The case of ``dg/classical_approach/multi_materials/solve_2d_magneto_static.py``
(same geometry, same constants), in continuous finite elements:

    -div(nu grad A) = -div(M)   in (0, 1)^2,    A = 0 on the boundary,

with the magnet ``(0.3, 0.5) x (0.3, 0.6)``: ``nu = 1 / 1.01`` and
``M = Bc e_y`` inside, ``nu = 1`` and ``M = 0`` outside.

Integrating by parts, the magnetization becomes a volume term,

    a(A, v) = int nu grad A . grad v,      l(v) = int M . grad v,

where DG has to carry it as the jump ``[[M]] . n`` on the interfaces. The
material is a named region of the mesh (``mesh.label_cells``) and the weak form
reads ``nu`` and ``M`` with ``per_region``: one read per cell, no coordinate
test. The grid resolves the magnet, so the region is exact.

Validation against the reference solution of ``Magneto_bi_validation.csv``
(the PINN examples' data): relative L2 error printed at the end.
"""

# %%
from pathlib import Path

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.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.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.abstract_weak_form import AbstractWeakForm

N_CELLS = 50  # as the DG example; the magnet edges fall on grid lines
ORDER = 1
MU_MAGNET = 1.01
BC = 1.0
MAGNET = (0.3, 0.5, 0.3, 0.6)  # x_min, x_max, y_min, y_max

CSV = (
    Path(__file__).parents[4]
    / "pinns/stationary_pdes/elliptic_pdes/Magneto_bi_validation.csv"
)


class MagnetoStaticWeakForm(AbstractWeakForm):
    """``nu grad A . grad v = M . grad v``, nu and M read per region."""

    def __init__(self, mu_magnet: float, bc: float):
        super().__init__(dim=2)
        self.mu_magnet = mu_magnet
        self.bc = bc

    def bilinear_form(self, u, v):
        nu = self.per_region({"magnet": 1.0 / self.mu_magnet}, default=1.0)
        return nu * u.gradient("x").dot(v.gradient("x"))

    def linear_form(self, v):
        m = self.per_region({"magnet": jnp.array([0.0, self.bc])}, default=jnp.zeros(2))
        return v.gradient("x").dot(m)


def identity(x):
    return x


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


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


def inside(box):
    x_min, x_max, y_min, y_max = box
    return lambda x: (
        (x[:, 0] > x_min) & (x[:, 0] < x_max) & (x[:, 1] > y_min) & (x[:, 1] < y_max)
    )


# %% Mesh, with the magnet as a named region.
mesh = Mesh(
    dim=2,
    n_cells=(N_CELLS, N_CELLS),
    ref_quad=UnitSquareTensorized(dim=2, order=ORDER + 2),
    mapping=Mapping(mappings=[InvertibleFunction(identity, identity)]),
)
mesh.label_cells("magnet", inside(MAGNET))

basis = AnalyticBasis(
    nb_basis=(ORDER + 1) ** 2,
    out_dim=1,
    mesh=mesh,
    local_basis=lagrange,
    basis_type="scalar",
)
scheme = EllipticFEscheme(
    AbstractPhysicalWeakModel.from_weak_form(
        MagnetoStaticWeakForm(MU_MAGNET, BC), dirichlet=zero
    ),
    VariablesFE(basis=basis, nb_variables=1),
)
scheme = EllipticFEscheme.solve(scheme)

# %% Validation against the reference solution.
data = np.genfromtxt(CSV, delimiter=";", skip_header=1)
points, reference = data[:, :2], data[:, 2]
u = np.asarray(scheme.variables.evaluate(jnp.asarray(points)))[:, 0]
error = u - reference
l2_rel = np.sqrt(np.mean(error**2)) / np.sqrt(np.mean(reference**2))
in_magnet = inside(MAGNET)(points)
print(
    f"CG-FEM Q{ORDER} {N_CELLS}x{N_CELLS}: relative L2 error vs reference "
    f"{l2_rel:.3e} (magnet {np.sqrt(np.mean(error[in_magnet] ** 2)):.2e}, "
    f"vacuum {np.sqrt(np.mean(error[~in_magnet] ** 2)):.2e} absolute)"
)

# %% Plot: solution, reference, error on the reference points.
fig, axes = plt.subplots(1, 3, figsize=(16, 4.8))
for ax, values, title, cmap in [
    (axes[0], u, "CG-FEM $A_h$", "jet"),
    (axes[1], reference, "reference", "jet"),
    (axes[2], np.abs(error), f"$|A_h - A|$, relative L2 {l2_rel:.2e}", "hot_r"),
]:
    scatter = ax.scatter(points[:, 0], points[:, 1], c=values, s=1, cmap=cmap)
    fig.colorbar(scatter, ax=ax)
    x_min, x_max, y_min, y_max = MAGNET
    ax.add_patch(
        plt.Rectangle((x_min, y_min), x_max - x_min, y_max - y_min, fill=False)
    )
    ax.set_title(title)
    ax.set_aspect("equal")
fig.tight_layout()
plt.show()