"""Lid-driven cavity in discontinuous Galerkin, Q2/Q1, against Ghia et al.

The DG counterpart of ``fem/.../solve_lid_driven_cavity_navier_stokes_2d.py``:
the same weak form (``NavierStokesWeakForm``), the faces of ``StokesDGFlux``
with the upwind convection term (``convection=True``, Lesaint-Raviart), the
lid imposed weakly through the flux -- so the corner singularity needs no
special treatment.

Solver: Newton, FGMRES, the block-triangular ``BlockSchurPreconditioner``
over the two spaces, velocity by one V-cycle of a DG multigrid (Schwarz
smoother, Galerkin coarse levels, relinearised at every Newton iterate),
Schur complement by LSC (``LSCInverse``): algebraic, it reads ``B``, ``B^T``
and ``F`` off the DG Jacobian's blocks, so nothing on the pressure space had
to be written for DG. Linear steps solved (``forcing="none"``), continuation
in Re; the viscosity is a pytree LEAF, so the whole continuation runs through
one compiled solve.

Measured (CPU; Newton / FGMRES / time with compilation, largest gap to
Ghia's u(0.5, y)):

=========  ============================  ============================
           Re = 100                      Re = 400 (from 100)
=========  ============================  ============================
16x16      5 / 166 / 11.2 s, 5.4e-3      5 / 543 / 12.8 s, 8.6e-3
32x32      4 / 187 / 29.3 s, 5.0e-3      5 / 790 / 88.3 s, 3.7e-3
=========  ============================  ============================

The same gaps as Taylor-Hood on the same meshes (5.6e-3 at 16x16, 5.0e-3 and
4.0e-3 at 32x32). Before the coarse levels were Galerkin (they were
rediscretised Laplacians, blind to the convection) and before a multigrid
relinearised inside the Newton step could read its blocks' triplets (it
probed its operator instead), 32x32 cost 4 / 663 / 86.7 s and
5 / 1 323 / 155.6 s.

Usage: ``python solve_lid_driven_cavity_navier_stokes_dg_2d.py [n_cells] [Re ...]
[--no-plot]``
(default 32, Re 100 then 400).
"""

import sys
import time

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

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.dg.elliptic_dg_scheme import (
    EllipticDGscheme,
)
from scimba_jax.linear_approximation.galerkin.dg.elliptic_dg_scheme_multi_space import (  # noqa: E501
    EllipticDGschemeMultipleSpaces,
)
from scimba_jax.linear_approximation.galerkin.dg.flux import SIPGFlux, StokesDGFlux
from scimba_jax.linear_approximation.meshes.mesh import Mesh
from scimba_jax.linear_approximation.quad.gauss_quad import UnitSquareTensorized
from scimba_jax.linear_approximation.solvers import NewtonKrylov
from scimba_jax.linear_approximation.solvers.multigrid import MG
from scimba_jax.linear_approximation.solvers.schur import (
    BlockSchurPreconditioner,
    CellBlockInverse,
    LSCInverse,
    MultigridInverse,
)
from scimba_jax.linear_approximation.solvers.smoothers import SchwarzSmoother
from scimba_jax.linear_approximation.transfer.hierarchy import (
    build_hierarchy_structured,
)
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.navier_stokes_weak_form import (
    NavierStokesWeakForm,
)
from scimba_jax.physical_models.classical_weakform.stokes_weak_form import (
    VectorLaplacianWeakForm,
)
from scimba_jax.physical_models.weak_boundary_conditions import Dirichlet
from scimba_jax.plots.plots_galerkin import plot_flow_2d

ORDER = 2  # velocity Q2, pressure Q1
SIGMA = 4.0 * ORDER * (ORDER + 1)
ARGS = [a for a in sys.argv[1:] if not a.startswith("-")]
N_CELLS = int(ARGS[0]) if ARGS else 32
REYNOLDS = tuple(float(r) for r in ARGS[1:]) or (100.0, 400.0)
N_COARSE = 4
U, P = 0, 1

# Ghia, Ghia & Shin (1982): u(0.5, y).
GHIA_Y = np.array([0.9766, 0.9688, 0.9609, 0.9531, 0.8516, 0.7344, 0.6172, 0.5,
                   0.4531, 0.2813, 0.1719, 0.1016, 0.0703, 0.0625, 0.0547])  # fmt: skip
GHIA_U = {
    100: [0.84123, 0.78871, 0.73722, 0.68717, 0.23151, 0.00332, -0.13641,
          -0.20581, -0.21090, -0.15662, -0.10150, -0.06434, -0.04775, -0.04192,
          -0.03717],
    400: [0.75837, 0.68439, 0.61756, 0.55892, 0.29093, 0.16256, 0.02135,
          -0.11477, -0.17119, -0.32726, -0.24299, -0.14612, -0.10338, -0.09266,
          -0.08186],
}  # fmt: skip


def make_mesh(n_cells):
    return Mesh(
        dim=2,
        n_cells=(n_cells, n_cells),
        ref_quad=UnitSquareTensorized(dim=2, order=ORDER + 2),
        mapping=Mapping(mappings=[InvertibleFunction(lambda x: x, lambda y: y)]),
    )


def make_space(mesh, order, out_dim, basis_type):
    basis = AnalyticBasis(
        nb_basis=(order + 1) ** 2,
        out_dim=out_dim,
        mesh=mesh,
        basis_type=basis_type,
        local_basis=lambda y, i, m: local_lagrange_basis(
            y, i, m, order=order, out_dim=out_dim
        ),
    )
    return VariablesDG(basis=basis, nb_variables=out_dim)


def lid(x):
    on_lid = (x[1] > 1.0 - 1e-12) & (x[0] > 1e-12) & (x[0] < 1.0 - 1e-12)
    return jnp.array([jnp.where(on_lid, 1.0, 0.0), 0.0])


def no_pressure_datum(x):
    """The pressure's boundary datum, which the flux never reads.

    ⚠ A module-level function, not a lambda in the Re loop: a fresh closure is
    a new object, it lands in the aux_data, and the solve would recompile at
    every Reynolds number although nothing it computes changed.
    """
    return jnp.zeros(1)


def velocity_multigrid():
    """A DG multigrid on the velocity space; its operator is read off J."""

    def scheme(n_cells):
        model = AbstractPhysicalWeakModel.from_weak_form(
            VectorLaplacianWeakForm(dim=2), dirichlet=lambda x: jnp.zeros(2)
        )
        return EllipticDGscheme(
            model,
            make_space(make_mesh(n_cells), ORDER, 2, "field"),
            SIPGFlux(sigma=SIGMA, h=1.0 / n_cells),
            use_scan_quad=False,
            linearization="blocks",
        )

    n_levels = int(np.log2(N_CELLS // N_COARSE)) + 1
    # Galerkin coarse levels (the builder's default): projected off the fine
    # block, they carry its convection, and a relinearisation traces no weak
    # form on them.
    hierarchy = build_hierarchy_structured(
        lambda counts: scheme(counts[0]), (N_CELLS, N_CELLS), n_levels
    )
    return MG(hierarchy, SchwarzSmoother())


def plot(velocity, pressure, re):
    """The PINN figure (streamlines, vorticity, pressure) and Ghia's u(0.5, y)."""
    fig, axes = plot_flow_2d(
        velocity,
        pressure,
        title=f"Lid-driven cavity, DG Q{ORDER}/Q{ORDER - 1} {N_CELLS}x{N_CELLS}, "
        f"Re = {re:g}",
        n_extra=1,
        show=False,
    )
    s = np.linspace(1e-6, 1.0 - 1e-6, 200)
    points = jnp.asarray(np.stack([0.5 * np.ones_like(s), s], -1))
    ax = axes[3]
    ax.plot(s, np.asarray(velocity.evaluate(points))[:, 0], "C0", label="u(0.5, s)")
    if int(re) in GHIA_U:
        ax.plot(GHIA_Y, GHIA_U[int(re)], "o", mfc="none", c="C0", label="Ghia et al.")
    ax.set_xlabel("s")
    ax.set_title("Vertical centreline")
    ax.grid(alpha=0.3)
    ax.legend()
    fig.tight_layout()


def main():
    mesh = make_mesh(N_CELLS)
    velocity = make_space(mesh, ORDER, 2, "field")
    pressure = make_space(mesh, ORDER - 1, 1, "scalar")
    mg = velocity_multigrid()
    # LSC's D: the diagonal of the (block-diagonal) DG velocity mass.
    mass = CellBlockInverse.mass_of(velocity)
    diagonal = 1.0 / np.asarray(
        np.diagonal(np.asarray(mass.inverse_blocks), axis1=1, axis2=2)
    ).reshape(-1)
    print(f"Lid-driven cavity, DG Q{ORDER}/Q{ORDER - 1}, {N_CELLS}x{N_CELLS}")
    dofsl = None
    for re in REYNOLDS:
        nu = 1.0 / re
        model = AbstractPhysicalWeakModel(dim=2)
        model.add_weak_form("main", NavierStokesWeakForm(dim=2, viscosity=nu))
        model.add_boundary_condition("0/boundary", Dirichlet(lid))
        model.add_boundary_condition("1/boundary", Dirichlet(no_pressure_datum))
        flux = StokesDGFlux(
            dim=2, viscosity=nu, sigma=SIGMA, h=1.0 / N_CELLS, convection=True
        )
        scheme = EllipticDGschemeMultipleSpaces(
            pde=model,
            variables_list=[velocity, pressure],
            flux_list=[flux, flux],
            equation_spaces=[0, 1],
            use_scan_quad=False,
            linearization="blocks",
        )
        pc = BlockSchurPreconditioner(
            sweep=[P, U],
            inverses={
                U: MultigridInverse(mg, component=U),
                P: LSCInverse(
                    jnp.asarray(diagonal), rtol=1e-4, max_iter=500, project_mean=True
                ),
            },
        )
        solver = NewtonKrylov(
            max_iter=25,
            tol=1e-8,
            cg_solver="fgmres:60",
            preconditioner=pc,
            max_iter_linear=600,
            rebuild_preconditioner=True,
            forcing="none",
        )
        started = time.perf_counter()
        dofsl, report = EllipticDGschemeMultipleSpaces.solve_pure(
            scheme, dofsl_init=dofsl, solver=solver
        )
        dofsl.block_until_ready()
        seconds = time.perf_counter() - started
        velocity.dofsl, pressure.dofsl = scheme._split_dofs(dofsl)
        points = jnp.stack([0.5 * jnp.ones(GHIA_Y.size), jnp.asarray(GHIA_Y)], -1)
        u_line = np.asarray(velocity.evaluate(points))[:, 0]
        gap = np.abs(u_line - np.array(GHIA_U[int(re)])).max()
        print(
            f"  Re {re:5.0f}: {dofsl.size} DOFs, Newton {int(report.n_iter)}, "
            f"FGMRES {int(report.n_linear)}, ||F|| {float(report.residual):.1e}, "
            f"{seconds:.1f} s | max |u - Ghia| {gap:.2e}",
            flush=True,
        )
        if "--no-plot" not in sys.argv:
            plot(velocity, pressure, re)
    plt.show()


if __name__ == "__main__":
    main()