"""2D magnetostatics with 5 magnets and 4 pole pieces, CG-FEM, named regions.

The N-material DG-SIPG case ported to continuous finite elements, same
geometry and constants:

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

* magnets ``(x0, x1) x (0.3, 0.6)`` for ``(x0, x1)`` in (0.3, 0.5), (0.6, 0.8),
  (0.9, 1.1), (1.2, 1.4), (1.5, 1.7): ``mu_r`` drawn in ``[1.005, 1.05]``,
  ``M = Bc e_y`` with ``Bc`` drawn in ``[0.4, 1.04]`` and its sign alternating
  from one magnet to the next;
* pole pieces ``(x0, x1) x (0.2, 0.7)`` for ``(x0, x1)`` in (0.5, 0.6),
  (0.8, 0.9), (1.1, 1.2), (1.4, 1.5): ``mu_r = 2000``;
* vacuum elsewhere, ``mu_r = 1``;

and ``nu = 1 / mu_r``. The weak form is the one of the two- and three-material
examples, ``a(A, v) = int nu grad A . grad v``, ``l(v) = int M . grad v``: each
magnet and pole piece is a named region of the mesh, and nine materials are
nine dictionary entries, not nine nested ``where``.

The draws are not seeded (as in the original case): each run is a new sample,
exported to ``csv_files/magnetostatic_mMpP_solution<i>.csv`` on a 201 x 201
grid.
"""

# %%
from pathlib import Path

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

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 = 80  # per direction on (0, 2) x (0, 1): every material edge on a grid line
ORDER = 1
QUAD_ORDER = 4
MU_POLAR = 2000.0
MAGNETS = [(0.3, 0.5), (0.6, 0.8), (0.9, 1.1), (1.2, 1.4), (1.5, 1.7)]  # x ranges
POLARS = [(0.5, 0.6), (0.8, 0.9), (1.1, 1.2), (1.4, 1.5)]
MAGNET_Y, POLAR_Y = (0.3, 0.6), (0.2, 0.7)


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

    Args:
        mu: ``{region: mu_r}``; regions not listed are vacuum.
        bc: ``{magnet region: Bc}``, the signed magnetization along ``e_y``.
    """

    def __init__(self, mu: dict, bc: dict):
        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(
            {k: jnp.array([0.0, b]) for k, b in self.bc.items()}, default=jnp.zeros(2)
        )
        return v.gradient("x").dot(m)


def stretch(p):
    return jnp.array([2.0 * p[0], p[1]])


def unstretch(p):
    return jnp.array([0.5 * p[0], p[1]])


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(x_range, y_range):
    (x0, x1), (y0, y1) = x_range, y_range
    return lambda x: (x[:, 0] > x0) & (x[:, 0] < x1) & (x[:, 1] > y0) & (x[:, 1] < y1)


# %% Materials: the draws of the original case, sign alternating along the row.
mu_m = np.random.uniform(1.005, 1.05, size=5)
bc_draw = np.random.uniform(0.4, 1.04, size=5)
bc = np.array([b if i % 2 == 0 else -b for i, b in enumerate(bc_draw)])
print("mu_m : ", mu_m)
print("Bc : ", bc)

mu = {f"magnet_{i}": float(m) for i, m in enumerate(mu_m, start=1)}
mu |= {f"polar_{i}": MU_POLAR for i in range(1, len(POLARS) + 1)}
magnetization = {f"magnet_{i}": float(b) for i, b in enumerate(bc, start=1)}

# %% Mesh of (0, 2) x (0, 1), the unit square stretched in x, with the regions.
mesh = Mesh(
    dim=2,
    n_cells=(N_CELLS, N_CELLS),
    ref_quad=UnitSquareTensorized(dim=2, order=QUAD_ORDER),
    mapping=Mapping(mappings=[InvertibleFunction(stretch, unstretch)]),
    is_identity_mapping=False,
)
for i, x_range in enumerate(MAGNETS, start=1):
    mesh.label_cells(f"magnet_{i}", inside(x_range, MAGNET_Y))
for i, x_range in enumerate(POLARS, start=1):
    mesh.label_cells(f"polar_{i}", inside(x_range, POLAR_Y))

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, magnetization), dirichlet=zero
    ),
    VariablesFE(basis=basis, nb_variables=1),
)
scheme = EllipticFEscheme.solve(scheme)

# %% Export on the 201 x 201 grid of the original case.
xs, ys = np.linspace(0.0, 2.0, 201), np.linspace(0.0, 1.0, 201)
XX, YY = np.meshgrid(xs, ys)
points = np.stack([XX.ravel(), YY.ravel()], axis=-1)
values = np.asarray(scheme.variables.evaluate(jnp.asarray(points)))[:, 0]

csv_dir = Path("csv_files")
csv_dir.mkdir(exist_ok=True)
index = 1
while (csv_dir / f"magnetostatic_mMpP_solution{index}.csv").exists():
    index += 1
csv_path = csv_dir / f"magnetostatic_mMpP_solution{index}.csv"
pd.DataFrame({"x": points[:, 0], "y": points[:, 1], "A": values}).to_csv(
    csv_path, index=False
)
print(f"CSV exported: {csv_path}")

# %% Plot.
fig, ax = plt.subplots(figsize=(10, 5))
field = values.reshape(XX.shape)
contour = ax.contourf(XX, YY, field, levels=40, cmap="jet")
ax.contour(XX, YY, field, levels=40, colors="k", linewidths=0.3)
fig.colorbar(contour, ax=ax)
for ranges, y_range, color in [(MAGNETS, MAGNET_Y, "white"), (POLARS, POLAR_Y, "cyan")]:
    for x0, x1 in ranges:
        ax.add_patch(
            plt.Rectangle(
                (x0, y_range[0]),
                x1 - x0,
                y_range[1] - y_range[0],
                fill=False,
                edgecolor=color,
                linewidth=1.5,
            )
        )
ax.set_title(f"CG-FEM Q{ORDER}, 5 magnets + 4 pole pieces, {N_CELLS}x{N_CELLS}")
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_aspect("equal")
fig.tight_layout()
plt.show()