"""Résolution d'un système 2D couplé d'advection-diffusion par DG.

Système sur Ω = [0, 1]² :
  -ε Δu + b₁·∇u + c₁·∇v = f   sur Ω,   u|∂Ω = 0
  -ε Δv + b₂·∇v + c₂·∇u = f   sur Ω,   v|∂Ω = 0

avec :
  ε      = 0.1   (diffusion)
  b₁     = (1, 0)    (advection propre de u)
  b₂     = (-1, 0)   (advection propre de v)
  c₁     = (0, 0.3)  (convection croisée : ∇v agit sur l'eq de u)
  c₂     = (0, -0.3) (convection croisée : ∇u agit sur l'eq de v)
  f(x,y) = Gaussienne centrée en (0.5, 0.5), σ = 0.1

Résolu comme un système vectoriel w = (u, v)ᵀ via DG (SIPG).
"""

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 import (
    EllipticDGscheme,
)
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

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

physical_dim = 2
out_dim = 2  # w = (u, v)ᵀ : deux inconnues scalaires
quad_order = 3  # points de Gauss par direction
poly_order = (
    1  # degré polynomial Q_1 → (1+1)² = 4 fonctions par cellule et par composante
)
n_cells = 40  # maillage n_cells × n_cells

sigma_gauss = 0.1  # largeur de la source gaussienne
x0 = jnp.array([0.5, 0.5])  # centre de la source


# ── Source gaussienne ─────────────────────────────────────────────────────────


def gaussian_source(x):
    """Source gaussienne G(x) = exp(-|x - x₀|² / (2σ²)) / (2πσ²)."""
    return jnp.exp(-jnp.sum((x - x0) ** 2) / (2.0 * sigma_gauss**2)) / (
        2.0 * jnp.pi * sigma_gauss**2
    )


# ── Forme faible du système ───────────────────────────────────────────────────


class SystemDiffAdvWeakForm(AbstractWeakForm):
    """Forme faible pour le système couplé d'advection-diffusion.

    Equation 1 (u) : -ε Δu + b₁·∇u + c₁·∇v = f
    Equation 2 (v) : -ε Δv + b₂·∇v + c₂·∇u = f
    """

    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.components()
        v_1, v_2 = v.components()

        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.components()
        f = self.get_fields()["f"]
        return ParamVecFunction.cat([f * v_1, f * v_2])


# ── Construction du problème ──────────────────────────────────────────────────

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,
)

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  # Q_1 → 4 fonctions par cellule

taylor_basis = AnalyticBasis(
    nb_basis=nb_basis,
    out_dim=out_dim,
    mesh=mesh,
    local_basis=lambda coords, i, m: local_taylor_basis(
        coords, i, m, order=poly_order, out_dim=out_dim
    ),
    basis_type="vec",  # variable vectorielle : (u, v)ᵀ
)

variables = VariablesDG(basis=taylor_basis, nb_variables=out_dim)

h = 1.0 / n_cells
sigma_sipg = poly_order * (poly_order + 1) * physical_dim
flux = SIPGFlux(sigma=sigma_sipg, h=h)


def dirichlet_bc(x):
    return jnp.zeros(out_dim)  # u = v = 0 sur ∂Ω


assembler = EllipticDGscheme(
    AbstractPhysicalWeakModel.from_weak_form(pde, dirichlet=dirichlet_bc),
    variables,
    flux,
)


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

print(f"Maillage : {n_cells}×{n_cells}, ordre p={poly_order}")
print(
    f"DOFs     : {assembler.variables.dofsl.size}  ({n_cells**2} × {nb_basis} × {out_dim})"
)
print("Résolution (Newton)…")

assembler = EllipticDGscheme.solve(
    assembler,
    max_iter=1,
    matrix_free=False,
    # tol=1e-5,
    # max_iter_linear=7000,
    # preconditioner=BlockJacobiPreconditioner(),
)
dofsl = assembler.variables.dofsl

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


# ── Évaluation sur une grille régulière ──────────────────────────────────────

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))  # (n_plot², 2)

w_h = np.array(assembler.variables.evaluate(pts))  # (n_plot², out_dim)
U = w_h[:, 0].reshape(n_plot, n_plot)
V = w_h[:, 1].reshape(n_plot, n_plot)

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

fig, axes = plt.subplots(1, 2, figsize=(10, 4))
fig.suptitle(
    f"Advection-diffusion DG — {n_cells}²×Q{poly_order}, σ_gauss={sigma_gauss}\n"
    f"b₁=(1,1)/√2  vs  b₂=(-1,-1)/√2",
    fontsize=11,
)

vmin, vmax = float(w_h.min()), float(w_h.max())

for ax, Z, title in zip(
    axes,
    [U, V],
    ["$u$ — advection $(1,1)/\\sqrt{2}$", "$v$ — advection $(-1,-1)/\\sqrt{2}$"],
):
    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()