"""Système couplé advection-diffusion avec deux espaces DG indépendants.

Même PDE que solve_2d_system_diffusion_advection.py :
  -ε Δu + b₁·∇u + c₁·∇v = f   sur Ω,   u|∂Ω = 0
  -ε Δv + b₂·∇v + c₂·∇u = f   sur Ω,   v|∂Ω = 0

Mais ici u et v sont deux ESPACES distincts (chacun sa base, son flux) sur le
même maillage. Le couplage croisé (c₁·∇v dans l'eq de u, etc.) est lu dans la
maille courante, comme n'importe quelle autre variable.

⚠ Une version antérieure mettait u et v sur deux maillages de finesses
différentes (30² et 20²), lus l'un depuis l'autre par ``find_cell_index``. Ce
chemin a été retiré (2026-09-11) : un maillage par domaine, une base par
variable -- voir ``check_same_refinement``. Les deux espaces sont ici sur le
maillage 30², le plus fin des deux.
"""

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

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_multi_space import (
    EllipticDGschemeMultipleSpaces,
)
from scimba_jax.linear_approximation.galerkin.dg.flux import SIPGFlux
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.nonlinear_approximation.model_class.funcparam_vectorial import (
    ParamVecFunction,
)
from scimba_jax.physical_models.abstract_physical_weak_model import (
    AbstractPhysicalWeakModel,
)
from scimba_jax.physical_models.abstract_weak_form import AbstractWeakForm
from scimba_jax.physical_models.weak_boundary_conditions import Dirichlet

# ── Paramètres ────────────────────────────────────────────────────────────────

physical_dim = 2
quad_order = 3
poly_order = 1

n_cells = 30

sigma_gauss = 0.1
x0 = jnp.array([0.5, 0.5])


def gaussian_source(x):
    return jnp.exp(-jnp.sum((x - x0) ** 2) / (2.0 * sigma_gauss**2)) / (
        2.0 * jnp.pi * sigma_gauss**2
    )


# ── Forme faible (identique au cas mono-space) ──────────────────────────────


class SystemDiffAdvWeakForm(AbstractWeakForm):
    def __init__(self, dim, eps, b1, b2, c1, c2, f):
        super().__init__(dim=dim)
        self.A = lambda x: eps * jnp.eye(dim)
        self.b1 = b1
        self.b2 = b2
        self.c1 = c1
        self.c2 = c2
        self.f = f

    def bilinear_form(
        self, u: ParamVecFunction, v: ParamVecFunction
    ) -> ParamVecFunction:
        u_1, u_2 = u  # un argument par espace
        v_1, v_2 = v

        grad_u1 = u_1.gradient("x")
        grad_u2 = u_2.gradient("x")
        grad_v1 = v_1.gradient("x")
        grad_v2 = v_2.gradient("x")

        fields = self.get_fields()
        A = fields["A"]
        b1 = fields["b1"]
        b2 = fields["b2"]
        c1 = fields["c1"]
        c2 = fields["c2"]

        diff_u = grad_v1.dot(A @ grad_u1)
        adv_u = grad_u1.dot(b1) * v_1
        cross_adv_u = grad_u2.dot(c1) * v_1

        diff_v = grad_v2.dot(A @ grad_u2)
        adv_v = grad_u2.dot(b2) * v_2
        cross_adv_v = grad_u1.dot(c2) * v_2

        return ParamVecFunction.cat(
            [diff_u + adv_u + cross_adv_u, diff_v + adv_v + cross_adv_v]
        )

    def linear_form(self, v: ParamVecFunction) -> ParamVecFunction:
        v_1, v_2 = v
        f = self.get_fields()["f"]
        return ParamVecFunction.cat([f * v_1, f * v_2])


# ── Construction du problème avec deux espaces ───────────────────────────────

b1_vec = jnp.array([1.0, 0.0])
b2_vec = jnp.array([-1.0, 0.0])
c1_vec = jnp.array([0.0, 0.3])
c2_vec = jnp.array([0.0, -0.3])

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

pde = SystemDiffAdvWeakForm(
    dim=physical_dim,
    eps=eps,
    b1=lambda x: b1_vec,
    b2=lambda x: b2_vec,
    c1=lambda x: c1_vec,
    c2=lambda x: c2_vec,
    f=gaussian_source,
)

# Un maillage, deux espaces (deux bases, deux flux)
mesh = Mesh(
    dim=physical_dim,
    n_cells=(n_cells, n_cells),
    ref_quad=UnitSquareTensorized(dim=physical_dim, order=quad_order),
    mapping=Mapping(mappings=[mapping_id]),
)
nb_basis = (poly_order + 1) ** physical_dim


def _basis():
    return AnalyticBasis(
        nb_basis=nb_basis,
        out_dim=1,
        mesh=mesh,
        local_basis=lambda coords, i, m: local_taylor_basis(
            coords, i, m, order=poly_order, out_dim=1
        ),
        basis_type="scalar",
    )


variables_u = VariablesDG(basis=_basis(), nb_variables=1)
variables_v = VariablesDG(basis=_basis(), nb_variables=1)

# Flux SIPG per-space
h = 1.0 / n_cells
sigma_sipg = poly_order * (poly_order + 1) * physical_dim
flux_u = SIPGFlux(sigma=sigma_sipg, h=h)
flux_v = SIPGFlux(sigma=sigma_sipg, h=h)


# Assembleur multi-space
def dirichlet_bc_u(x):
    return jnp.zeros(1)


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


# Le modèle porte la physique ET les CL : une par espace, dans l'ordre des espaces.
model = AbstractPhysicalWeakModel.from_weak_form(pde)
model.add_boundary_condition("0/boundary", Dirichlet(dirichlet_bc_u))
model.add_boundary_condition("1/boundary", Dirichlet(dirichlet_bc_v))

assembler = EllipticDGschemeMultipleSpaces(
    pde=model,
    variables_list=[variables_u, variables_v],
    flux_list=[flux_u, flux_v],
    equation_spaces=[0, 1],
)


# ── Résolution ────────────────────────────────────────────────────────────────

ndof_u = variables_u.ndof_linear
ndof_v = variables_v.ndof_linear
print(f"Espace u : {n_cells}²×Q{poly_order}, DOFs = {ndof_u}")
print(f"Espace v : {n_cells}²×Q{poly_order}, DOFs = {ndof_v}")
print(f"Total DOFs : {ndof_u + ndof_v}")
print("Résolution (Newton)…")

assembler = EllipticDGschemeMultipleSpaces.solve(
    assembler, max_iter=1, matrix_free=False
)

print("Résolution terminée.")


# ── Évaluation sur une grille commune ────────────────────────────────────────

n_plot = 60
xp = np.linspace(0.0, 1.0, n_plot)
XC, YC = np.meshgrid(xp, xp)
pts = jnp.array(np.stack([XC.ravel(), YC.ravel()], axis=-1))

U = np.array(assembler.variables_list[0].evaluate(pts)).reshape(n_plot, n_plot)
V = np.array(assembler.variables_list[1].evaluate(pts)).reshape(n_plot, n_plot)

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

fig, axes = plt.subplots(1, 2, figsize=(10, 4))
fig.suptitle(
    f"Multi-space DG — u, v : {n_cells}² ×Q{poly_order}\nCross-convection coupling",
    fontsize=11,
)

w_h = np.concatenate([U.ravel()[:, None], V.ravel()[:, None]], axis=-1)
vmin, vmax = float(w_h.min()), float(w_h.max())

for ax, Z, title in zip(
    axes,
    [U, V],
    ["$u$ — advection $(1,0)$", "$v$ — advection $(-1,0)$"],
):
    im = ax.pcolormesh(XC, YC, Z, shading="auto", cmap="turbo", vmin=vmin, vmax=vmax)
    ax.contour(XC, YC, Z, levels=10, colors="k", linewidths=0.5, alpha=0.4)
    ax.set_title(title)
    ax.set_aspect("equal")
    plt.colorbar(im, ax=ax)

plt.tight_layout()
plt.show()