"""2D magnetostatics with a magnet and a pole piece, CG-FEM, three named regions.

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

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

* magnet ``(0.3, 0.5) x (0.3, 0.6)``: ``mu_r = 1.01``, ``M = Bc e_y``;
* pole piece ``(0.5, 0.6) x (0.2, 0.7)``: ``mu_r = 2000``, ``M = 0``;
* vacuum elsewhere: ``mu_r = 1``, ``M = 0``;

and ``nu = 1 / mu_r``. As in the two-material example, ``l(v) = int M . grad v``
and both coefficients are read per region: adding the pole piece is one more
``label_cells`` and one more dictionary entry, not a nested ``where``.

Validation against the reference solution of ``Magneto_tri_validation.csv``.
"""

# %%
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 = 60  # as the DG example; every material edge falls on a grid line
ORDER = 1
MU = {"magnet": 1.01, "polar": 2000.0}
BC = 1.0
BOXES = {  # x_min, x_max, y_min, y_max
    "magnet": (0.3, 0.5, 0.3, 0.6),
    "polar": (0.5, 0.6, 0.2, 0.7),
}

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


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

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

    def bilinear_form(self, u, v):
        nu = self.per_region({k: 1.0 / mu for k, mu in self.mu.items()}, 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 and the pole piece as named regions.
mesh = Mesh(
    dim=2,
    n_cells=(N_CELLS, N_CELLS),
    ref_quad=UnitSquareTensorized(dim=2, order=ORDER + 2),
    mapping=Mapping(mappings=[InvertibleFunction(identity, identity)]),
)
for name, box in BOXES.items():
    mesh.label_cells(name, inside(box))

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, 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))
masks = {name: inside(box)(points) for name, box in BOXES.items()}
vacuum = ~(masks["magnet"] | masks["polar"])
print(
    f"CG-FEM Q{ORDER} {N_CELLS}x{N_CELLS}: relative L2 error vs reference "
    f"{l2_rel:.3e} (absolute: vacuum {np.sqrt(np.mean(error[vacuum] ** 2)):.2e}, "
    + ", ".join(
        f"{name} {np.sqrt(np.mean(error[mask] ** 2)):.2e}"
        for name, mask in masks.items()
    )
    + ")"
)

# %% 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)
    for x_min, x_max, y_min, y_max in BOXES.values():
        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()