r"""Lid-driven cavity flow (stationary Navier-Stokes, Re=100).

Geometry: unit square [0, 1]².

Equations (isothermal incompressible NS):

.. math::

    (u \cdot \nabla)u - \nu \Delta u + \nabla p &= 0 \\
    \nabla \cdot u &= 0

Boundary conditions:

* north (y=1):  u = (sin²(πx), 0)  — regularised lid — **hard constraint**
* south (y=0):  u = (0, 0)         — no-slip         — **hard constraint**
* west/east:    u = (0, 0)         — no-slip         — **hard constraint**

The standard u=1 lid creates a singularity at the top corners (u must
simultaneously be 1 on the north wall and 0 on the side walls).  The
regularised condition u_lid = sin²(πx) vanishes smoothly at both corners,
so all four BCs can be enforced as hard constraints without any singularity.
The resulting vortex is physically identical to the u=1 case in the interior.

Hard constraint via post-processing:

* u  → u_net · φ + sin²(πx) · y,  φ = x(1−x)y(1−y)
       → u=0 on south/west/east,  u=sin²(πx) on north
* v  → v_net · φ  → v=0 on all walls
* p  → unconstrained (determined up to a constant)

No BC residuals are needed: all boundary conditions are satisfied exactly.

A single ``MODE`` flag selects the approximation-space layout:

* ``"vec"``          — one network, output ``[u_x, u_y, p]``
* ``"field_scalar"`` — two networks: one field network for ``u``,
                       one scalar network for ``p``

Uses :class:`NavierStokesND` from
``scimba_jax.physical_models.elliptic_pde.navier_stokes_stationary``.
"""

import timeit

import jax
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.nonlinear_approximation.approximation_spaces.approximation_spaces import (
    ApproximationSpace,
)
from scimba_jax.nonlinear_approximation.integration.monte_carlo import (
    DomainSampler,
    TensorizedSampler,
)
from scimba_jax.nonlinear_approximation.model_class.funcparam_vectorial import (
    ParamVecFunction,
)
from scimba_jax.nonlinear_approximation.networks.mlp import MLP
from scimba_jax.nonlinear_approximation.numerical_solvers.projectors import Projector
from scimba_jax.physical_models.elliptic_pde.navier_stokes_stationary import (
    NavierStokesND,
    _extract_u_p,
)
from scimba_jax.plots.plots_nd import plot_abstract_approx_space

# ── configuration ─────────────────────────────────────────────────────────────

# "vec"          → one network, output [u_x, u_y, p]
# "field_scalar" → two networks: field net for u, scalar net for p
MODE = "vec"

# ── hyper-parameters ──────────────────────────────────────────────────────────

N_COLLOC = 7000
N_EPOCHS = 500

# Re = u_ref * L / nu = 1 * 1 / 0.0005 = 2000.0
nu = 0.0005

# Loss weights — interior only: [NS_x, NS_y, div]
WEIGHTS = {
    "interior": [1.0, 1.0, 2.0],
}

# ── hard-constraint post-processing ──────────────────────────────────────────
# u_lift = sin²(πx) · y  satisfies u=0 on all four walls and u=sin²(πx) at
# the lid (y=1).  Unlike the naive lift u=y, this lift does NOT satisfy the
# NS interior equations, so the network cannot escape into the trivial
# Couette solution.


def post_processing_uv(output, x):
    x_, y_ = x[0], x[1]
    phi = x_ * (1 - x_) * y_ * (1 - y_)
    u_lift = jnp.sin(jnp.pi * x_) ** 2 * y_
    return jnp.array([output[0] * phi + u_lift, output[1] * phi])


def post_processing_p(output, x):
    return output  # pressure: unconstrained


def post_processing_vec(output, x):
    x_, y_ = x[0], x[1]
    phi = x_ * (1 - x_) * y_ * (1 - y_)
    u_lift = jnp.sin(jnp.pi * x_) ** 2 * y_
    u = output[0] * phi + u_lift
    v = output[1] * phi
    p = output[2]
    return jnp.array([u, v, p])


# ── domain ────────────────────────────────────────────────────────────────────

key = jax.random.PRNGKey(0)

box = [(0.0, 1.0), (0.0, 1.0)]

dx = Square2D(box, is_main_domain=True)

# ── physics model — no BC residuals (all BCs are hard) ───────────────────────

model = NavierStokesND(dx, dim=2, nu=nu, mode=MODE, bc_residuals={})

# ── approximation space (depends on MODE) ─────────────────────────────────────

if MODE == "vec":
    nn = MLP(in_size=2, out_size=3, hidden_sizes=[16, 16, 16, 16], key=key)
    space = ApproximationSpace(
        {"x": 2},
        [(nn, "vec", 3)],
        model_type="x",
        post_processing=post_processing_vec,
    )

else:  # "field_scalar"
    nn_uv = MLP(in_size=2, out_size=2, hidden_sizes=[18, 18, 18, 18], key=key)
    nn_p = MLP(in_size=2, out_size=1, hidden_sizes=[16, 16, 16], key=key)
    space = ApproximationSpace(
        {"x": 2},
        [(nn_uv, "field", 2), (nn_p, "scalar", 1)],
        model_type="x",
        post_processing=[post_processing_uv, post_processing_p],
    )

# ── sampler ───────────────────────────────────────────────────────────────────

sampler = TensorizedSampler(
    [DomainSampler(dx)],
    bc=False,
)

# ── training ──────────────────────────────────────────────────────────────────

pinn = Projector(model, space, sampler, weights=WEIGHTS, one_loss_by_equation=True)

key, sample_dict = sampler.sample(key, N_COLLOC, 0)
loss0 = pinn.evaluate_loss(space, sample_dict)
print("initial loss:", loss0)

start = timeit.default_timer()
key, pinn = pinn.project(key, space, N_EPOCHS, N_COLLOC, 0)
end = timeit.default_timer()

print(f"best loss: {pinn.best_loss}")
print(f"time for {N_EPOCHS} epochs: {end - start:.1f}s")

# ── visualisation ─────────────────────────────────────────────────────────────

plot_abstract_approx_space(
    pinn.space,
    dx,
    components=[0, 1, 2],
    loss=pinn.losses,
    residual=pinn.model,
    draw_contours=True,
    n_drawn_contours=20,
)

# ── streamlines + vorticity (JAX autodiff, classical cavity style) ────────────

N_GRID = 100
x_lin = np.linspace(0.0, 1.0, N_GRID)
y_lin = np.linspace(0.0, 1.0, N_GRID)
X, Y = np.meshgrid(x_lin, y_lin)
xy_flat = np.stack([X.ravel(), Y.ravel()], axis=1)
xt = jnp.array(xy_flat)

# Build symbolic functions — vorticity via curl_x() (exact JAX autodiff)
variables = pinn.space.create_variables()
all_together = ParamVecFunction.cat(variables)
u_field, _ = _extract_u_p(variables, MODE, 2)
omega_func = u_field.curl_x()  # ∂v/∂x − ∂u/∂y (scalar in 2D)

# Evaluate on grid
w_pred = jax.device_get(all_together.vmap_on_physical_variables()(pinn.space, xt))
omega_vals = jax.device_get(omega_func.vmap_on_physical_variables()(pinn.space, xt))

U = w_pred[:, 0].reshape(N_GRID, N_GRID)
V = w_pred[:, 1].reshape(N_GRID, N_GRID)
P = w_pred[:, 2].reshape(N_GRID, N_GRID)
omega = omega_vals.reshape(N_GRID, N_GRID)

Re = int(round(1.0 / nu))

fig, axes = plt.subplots(1, 3, figsize=(16, 5))
fig.suptitle(f"Lid-driven cavity  Re={Re}", fontsize=13)

# — streamlines coloured by speed —
ax = axes[0]
speed = np.sqrt(U**2 + V**2)
strm = ax.streamplot(X, Y, U, V, color=speed, cmap="RdBu_r", density=2, linewidth=0.8)
fig.colorbar(strm.lines, ax=ax, label="|u|")
ax.set_title("Streamlines")
ax.set_aspect("equal")
ax.set_xlabel("x")
ax.set_ylabel("y")

# — vorticity —
ax = axes[1]
om_lim = np.percentile(np.abs(omega), 99)
levels_om = np.linspace(-om_lim, om_lim, 40)
cf = ax.contourf(X, Y, omega, levels=levels_om, cmap="RdBu_r", extend="both")
ax.contour(X, Y, omega, levels=levels_om[::4], colors="k", linewidths=0.4)
fig.colorbar(cf, ax=ax, label="ω")
ax.set_title("Vorticity  ω = ∂v/∂x − ∂u/∂y")
ax.set_aspect("equal")
ax.set_xlabel("x")
ax.set_ylabel("y")

# — pressure —
ax = axes[2]
cf = ax.contourf(X, Y, P, levels=40, cmap="RdBu_r")
ax.contour(X, Y, P, levels=20, colors="k", linewidths=0.4)
fig.colorbar(cf, ax=ax, label="p")
ax.set_title("Pressure p")
ax.set_aspect("equal")
ax.set_xlabel("x")
ax.set_ylabel("y")

plt.tight_layout()
plt.show()