"""Magnetostatic DG-SIPG à 3 matériaux.

Domaines :
  Ω_v  vacuum  : (0,1)² \\ (Ω_m ∪ Ω_p),  μ_r = 1
  Ω_m  magnet  : (0.3, 0.5) × (0.3, 0.6), μ_r = 1.01,  M = Bc·ê_y
  Ω_p  polar   : (0.5, 0.6) × (0.2, 0.7), μ_r = 2000
"""

from pathlib import Path

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

from scimba_jax.domains.meshless_domains.domains_2d import Square2D
from scimba_jax.linear_approximation.basis.analytic_bases import local_taylor_basis
from scimba_jax.linear_approximation.basis.general_bases import AnalyticBasis
from scimba_jax.linear_approximation.galerkin.dg.elliptic_dg_scheme import (
    EllipticDGscheme,
)
from scimba_jax.linear_approximation.galerkin.dg.flux import MagnetoStaticSIPGFlux
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_dg import VariablesDG
from scimba_jax.mapping.mapping import InvertibleFunction, Mapping
from scimba_jax.physical_models.abstract_physical_weak_model import (
    AbstractPhysicalWeakModel,
)
from scimba_jax.physical_models.classical_weakform.diffusion_advection_reaction_weak_form import (
    EllipticWeakForm,
)

# ── Weak form 3 matériaux ──────────────────────────────────────────────────────


class MagnetoStaticWeakForm3Mat(EllipticWeakForm):
    """Weak form -div(A(x) ∇u) = -div(M(x)) pour 3 sous-domaines.

    A(x) = 1/μ_r(x) avec :
      μ_r = 1      dans le vacuum
      μ_r = mu_m   dans l'aimant
      μ_r = mu_p   dans la pièce polaire

    M(x) = Bc·ê_y dans l'aimant, 0 ailleurs.
    Exposé comme ``self.M`` pour que ``MagnetoStaticSIPGFlux`` ajoute
    le terme de surface [[M·n]] à chaque interface.
    """

    def __init__(
        self,
        magnet_domain: Square2D = None,
        polar_domain: Square2D = None,
        Bc: float = 1.0,
        mu_m: float = 1.01,
        mu_p: float = 2000.0,
    ):
        dim = 2
        self.magnet_domain = magnet_domain
        self.polar_domain = polar_domain
        self.Bc = Bc

        is_mag = magnet_domain.is_inside
        is_pol = polar_domain.is_inside

        super().__init__(
            dim=dim,
            A=lambda x: jnp.eye(dim)
            * jnp.where(
                is_mag(x[None])[0],
                1.0 / mu_m,
                jnp.where(is_pol(x[None])[0], 1.0 / mu_p, 1.0),
            ),
            b=lambda x: jnp.zeros(dim),
            c=lambda x: jnp.zeros(()),
            f=lambda x: jnp.zeros(()),
        )

        self.M = lambda x: jnp.where(
            is_mag(x[None])[0], jnp.array([0.0, Bc]), jnp.zeros(2)
        )

    def bilinear_form(self, u, v):
        fields = self.get_fields()
        return v.gradient("x").dot(fields["A"] @ u.gradient("x"))

    def linear_form(self, v):
        return self.get_fields()["f"] * v


# ── Helpers visuels ───────────────────────────────────────────────────────────


def add_rects(ax):
    for x0, x1, y0, y1, color, label in [
        (0.3, 0.5, 0.3, 0.6, "white", r"$\Omega_m$"),
        (0.5, 0.6, 0.2, 0.7, "cyan", r"$\Omega_p$"),
    ]:
        rect = plt.Rectangle(
            (x0, y0),
            x1 - x0,
            y1 - y0,
            fill=False,
            edgecolor=color,
            linewidth=1.5,
            label=label,
        )
        ax.add_patch(rect)


def plot_solution_2d(assembler, n_cells_val, n_plot=201):
    xs = np.linspace(0.0, 1.0, n_plot)
    ys = np.linspace(0.0, 1.0, n_plot)
    XX, YY = np.meshgrid(xs, ys)
    pts = jnp.array(np.stack([XX.ravel(), YY.ravel()], axis=-1))

    UH = np.array(assembler.variables.evaluate(pts))[:, 0].reshape(n_plot, n_plot)

    LEVELS = 40
    fig, ax = plt.subplots(figsize=(6, 5))
    cf = ax.contourf(XX, YY, UH, levels=LEVELS, cmap="jet")
    ax.contour(XX, YY, UH, levels=LEVELS, colors="k", linewidths=0.3)
    fig.colorbar(cf, ax=ax)
    add_rects(ax)
    ax.legend(loc="upper right", fontsize=9)
    ax.set_title(f"DG SIPG 3 mat. — {n_cells_val}×{n_cells_val} — $u_h$")
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_aspect("equal")
    plt.tight_layout()
    plt.show()


# ── Domaines ──────────────────────────────────────────────────────────────────

square_domain = Square2D(bounds=[(0.0, 1.0), (0.0, 1.0)], is_main_domain=True)
magnet_domain = Square2D(bounds=[(0.3, 0.5), (0.3, 0.6)], is_main_domain=False)
polar_domain = Square2D(bounds=[(0.5, 0.6), (0.2, 0.7)], is_main_domain=False)

# ── PDE ───────────────────────────────────────────────────────────────────────

pde = MagnetoStaticWeakForm3Mat(
    magnet_domain=magnet_domain,
    polar_domain=polar_domain,
)

# ── Maillage ──────────────────────────────────────────────────────────────────

physical_dim = 2
out_dim = 1
order = 1
nb_basis = (order + 1) ** physical_dim
quad_order = 4
n_cells = 60
h_phys = 1.0 / n_cells

mapping_id = InvertibleFunction(lambda x: x, lambda y: y)

m = Mesh(
    dim=physical_dim,
    n_cells=(n_cells, n_cells),
    ref_quad=UnitSquareTensorized(dim=physical_dim, order=quad_order),
    mapping=Mapping(mappings=[mapping_id]),
    is_identity_mapping=False,
)


def dirichlet_bc(x):
    return jnp.zeros(out_dim)


# ── Solveur ───────────────────────────────────────────────────────────────────

taylor_basis = AnalyticBasis(
    nb_basis=nb_basis,
    out_dim=out_dim,
    mesh=m,
    local_basis=lambda coords, i, mesh: local_taylor_basis(
        coords, i, mesh, order=order, out_dim=out_dim
    ),
    basis_type="scalar",
)
flux = MagnetoStaticSIPGFlux(sigma=order * (order + 1), h=h_phys)
variables = VariablesDG(basis=taylor_basis, nb_variables=out_dim)
assembler = EllipticDGscheme(
    AbstractPhysicalWeakModel.from_weak_form(pde, dirichlet=dirichlet_bc),
    variables,
    flux,
)

assembler = EllipticDGscheme.solve(assembler, max_iter=1, matrix_free=True, tol=1e-7)

# ── Visualisation ─────────────────────────────────────────────────────────────

print("\nPlot solution...")
plot_solution_2d(assembler, n_cells, n_plot=201)

# ── Validation vs CSV ─────────────────────────────────────────────────────────

csv_path = (
    Path(__file__).parent.parent.parent.parent
    / "pinns/stationary_pdes/elliptic_pdes/Magneto_tri_validation.csv"
)

if csv_path.exists():
    data_val = np.genfromtxt(csv_path, delimiter=";", skip_header=1)
    x_val = jnp.array(data_val[:, 0])
    y_val = jnp.array(data_val[:, 1])
    A_ref = jnp.array(data_val[:, 2])
    xy_val = jnp.stack([x_val, y_val], axis=-1)

    u_val = np.array(assembler.variables.evaluate(xy_val))[:, 0]
    A_ref_np = np.array(A_ref)
    x_val_np = np.array(x_val)
    y_val_np = np.array(y_val)

    err = u_val - A_ref_np
    err_abs = np.abs(err)
    l2_err = float(np.sqrt(np.mean(err**2)))
    l2_ref = float(np.sqrt(np.mean(A_ref_np**2)))
    l2_rel = l2_err / l2_ref
    linf_err = float(np.max(err_abs))
    linf_rel = linf_err / float(np.max(np.abs(A_ref_np)))

    mask_m = (
        (x_val_np >= 0.3) & (x_val_np <= 0.5) & (y_val_np >= 0.3) & (y_val_np <= 0.6)
    )
    mask_p = (
        (x_val_np >= 0.5) & (x_val_np <= 0.6) & (y_val_np >= 0.2) & (y_val_np <= 0.7)
    )

    l2_m = float(np.sqrt(np.mean(err[mask_m] ** 2))) if mask_m.any() else float("nan")
    l2_p = float(np.sqrt(np.mean(err[mask_p] ** 2))) if mask_p.any() else float("nan")
    l2_v = float(np.sqrt(np.mean(err[~mask_m & ~mask_p] ** 2)))

    print(
        f"\nValidation CSV → L2={l2_err:.3e}  L2_rel={l2_rel:.3e}  "
        f"Linf={linf_err:.3e}  Linf_rel={linf_rel:.3e}"
    )
    print(f"               L2 vacuum={l2_v:.3e}  magnet={l2_m:.3e}  polar={l2_p:.3e}")

    sort_idx = np.argsort(x_val_np)
    x_s = x_val_np[sort_idx]
    y_s = y_val_np[sort_idx]
    A_s = A_ref_np[sort_idx]
    u_s = u_val[sort_idx]
    err_s = err_abs[sort_idx]

    fig3, axs3 = plt.subplots(1, 3, figsize=(18, 5))

    sc0 = axs3[0].scatter(x_s, y_s, c=A_s, cmap="jet", s=1)
    fig3.colorbar(sc0, ax=axs3[0])
    axs3[0].set_title("Référence CSV")

    sc1 = axs3[1].scatter(
        x_s, y_s, c=u_s, cmap="jet", s=1, vmin=A_s.min(), vmax=A_s.max()
    )
    fig3.colorbar(sc1, ax=axs3[1])
    axs3[1].set_title("DG SIPG 3 mat.")

    vmax_err = float(np.percentile(err_abs, 99))
    sc2 = axs3[2].scatter(x_s, y_s, c=err_s, cmap="hot_r", s=1, vmin=0, vmax=vmax_err)
    fig3.colorbar(sc2, ax=axs3[2])
    axs3[2].set_title(f"|DG − ref|  L2={l2_err:.2e}  L2_rel={l2_rel:.2e}")

    for ax in axs3:
        ax.set_xlabel("x")
        ax.set_ylabel("y")
        ax.set_aspect("equal")
        add_rects(ax)

    plt.suptitle(f"DG SIPG {n_cells}×{n_cells} 3 matériaux vs CSV", fontsize=12)
    plt.tight_layout()
    plt.show()
else:
    print(f"\nCSV non trouvé : {csv_path}")