"""Steady Euler flow around a NACA0012 airfoil, first-order FV on a Gmsh mesh.

    div F(W) = 0  around the airfoil,  W = (rho, rho v, E),  gamma = 1.4
    M_inf = 0.8, alpha = 1.25 deg  (AGARD-AR-211 test case 01, transonic)

The classical airfoil case: a strong shock on the upper surface near
x/c = 0.6 and a weak one on the lower surface near 0.35, lift and drag
against the fine-grid references (Vassberg & Jameson 2010: C_L = 0.3541,
C_D = 0.0223). A first-order scheme does not reach them -- its numerical
dissipation shows as a smeared shock and as spurious drag -- but it moves
towards them under refinement: measured C_L = 0.252 / C_D = 0.049 on 5 700
cells (h = 0.01c at the wall, 8 s), 0.310 / 0.034 on 21 700 (h = 0.005c,
21 s); the shocks sit where they should and C_p,max reaches 1.147 for an
isentropic 1.170. Set
``MACH, ALPHA_DEG = 0.5, 0.0`` for the subsonic symmetric case (C_L = C_D = 0
exactly in the continuous problem: the residual C_D is the dissipation).

What the file needed, and now has:

* a Gmsh mesh with a hole whose boundary is smooth **by pieces** -- upper
  and lower surfaces two splines meeting at a sharp trailing edge
  (``_add_loop`` on a list of pieces) -- and a size **graded** by the
  distance to the airfoil (``mesh_size`` as a callable), 0.01c at the wall
  to 1.5c at the far-field circle of radius 15c (0.005c by default);
* named boundaries on an unstructured mesh, ``mesh.label_boundary``, so the
  ``Wall`` (slip: mirror state) goes on the airfoil and the free stream on
  the far field, keyed by name as everywhere;
* a pseudo-time march with a local time step per cell
  (``march_to_steady_state``): with ``h`` spanning two decades a global
  explicit step would need ~1e5 iterations, the local one a few thousand.

The problem is STEADY, and the march IS the solver: forward Euler on
``W_t + div F(W) = 0`` from the uniform free stream, every cell at its own
CFL step, until the density residual has dropped by five decades. The
intermediate states are not a flow at any time; only the fixed point means
something, and it is the same whatever the step.

Newton-Krylov on the same FV residual (``FiniteVolumeScheme.solve``,
matrix-free JVP, GMRES) is then run as a finish from the marched state, to
see what it buys. Measured on 5 700 cells, with the linear solve converged (restart 100), with or without
Armijo or Eisenstat-Walker forcing, five steps reduce |F| by x0.41 from a
rough state and by x0.89 from a converged one -- the Newton direction of a
first-order upwind residual that is only piecewise smooth (the ``where``
branches of HLLC, the min/max of the Davis speeds, the wall mirror) buys
what a descent step buys. Newton needs a smoother residual, i.e. a viscous
term, to pay off here; the local pseudo-time march, a Newton damped by the
diagonal, is what converges this inviscid first-order problem.

The fluxes are the n-D solvers of ``euler_sod_1d.py`` and the cylindrical
explosion, unchanged: an unstructured face is a normal like any other.

Needs the optional ``mesh`` extra (gmsh).
"""

import time

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.collections import PolyCollection
from matplotlib.path import Path
from matplotlib.tri import Triangulation

from scimba_jax.linear_approximation.finite_volume import (
    FiniteVolumeScheme,
    march_to_steady_state,
)
from scimba_jax.linear_approximation.finite_volume.hyperbolic import EulerHLLCFlux
from scimba_jax.linear_approximation.meshes.unstructured_mesh import UnstructuredMesh
from scimba_jax.linear_approximation.quad.gauss_quad import UnitSquareTensorized
from scimba_jax.linear_approximation.solvers.newton import NewtonKrylov
from scimba_jax.linear_approximation.variables.variables_fv import VariablesFV
from scimba_jax.mapping.macro_mesh import ExactCircle, macro_mesh_from_loops
from scimba_jax.physical_models.abstract_physical_conservative_model import (
    AbstractPhysicalConservativeModel,
)
from scimba_jax.physical_models.classical_weakform.euler_weak_form import (
    EulerPDE,
    euler_conserved,
    euler_primitives,
)
from scimba_jax.physical_models.weak_boundary_conditions import Dirichlet, Wall

GAMMA = 1.4
MACH, ALPHA_DEG = 0.8, 1.25  # AGARD 01; (0.5, 0.0) for the subsonic symmetric case
FARFIELD_RADIUS = 15.0  # chords, centred on mid-chord
H_WALL, H_FAR, GROWTH = 0.003, 1.5, 0.03  # graded size: h = H_WALL + GROWTH * distance
N_AIRFOIL_POINTS = 160  # per surface, cosine-clustered
CFL = 0.8
N_STEPS, REPORT_EVERY = 30000, 2000
# The Newton-Krylov finish, from the marched state: NEWTON_ITER steps, each
# solved by GMRES with a restart of 100 and KRYLOV_CYCLES restart cycles.
# ⚠ ``max_iter_linear`` counts restart cycles for jax's GMRES, not matvecs;
# and the restart matters: at 30 or 40 GMRES stagnates on this Jacobian
# (relative linear residual 0.95 per cycle, measured), at 100 it converges
# (5e-2 in one cycle, 2e-5 in four).
NEWTON_ITER, KRYLOV_CYCLES, GMRES_RESTART = 10, 2, 100

# Free stream, non-dimensional: rho = 1, c = 1, so |v| = M and p = 1 / gamma.
ALPHA = np.deg2rad(ALPHA_DEG)
V_INF = MACH * np.array([np.cos(ALPHA), np.sin(ALPHA)])
P_INF = 1.0 / GAMMA
W_INF = euler_conserved(1.0, jnp.asarray(V_INF), P_INF, GAMMA)
DYNAMIC_PRESSURE = 0.5 * MACH * MACH  # 1/2 rho_inf |v_inf|^2


# ── Geometry and mesh ────────────────────────────────────────────────────────
def naca0012_surfaces(n=N_AIRFOIL_POINTS):
    """Upper surface TE -> LE and lower surface LE -> TE, closed trailing edge.

    The 4-digit thickness law with the last coefficient -0.1036 (instead of
    -0.1015) so that the trailing edge closes exactly at (1, 0). Cosine
    clustering puts the points where the curvature is, at the leading edge.
    """
    beta = np.linspace(0.0, np.pi, n)
    x = 0.5 * (1.0 - np.cos(beta))
    thickness = 0.6 * (
        0.2969 * np.sqrt(x) - 0.1260 * x - 0.3516 * x**2 + 0.2843 * x**3 - 0.1036 * x**4
    )
    upper = np.stack([x[::-1], thickness[::-1]], axis=1)
    lower = np.stack([x, -thickness], axis=1)
    return [upper, lower]


def make_mesh():
    surfaces = naca0012_surfaces()
    airfoil_points = np.vstack(surfaces)

    def mesh_size(x, y):
        distance = np.min(np.hypot(airfoil_points[:, 0] - x, airfoil_points[:, 1] - y))
        return float(min(H_WALL + GROWTH * distance, H_FAR))

    macro = macro_mesh_from_loops(
        ExactCircle(0.5, 0.0, FARFIELD_RADIUS),
        holes=[surfaces],
        order=1,
        mesh_size=mesh_size,
        smooth=True,
    )
    mesh = UnstructuredMesh.from_macro_mesh(macro, UnitSquareTensorized(dim=2, order=2))
    # The two boundaries, told apart by where they are: the airfoil sits
    # within a chord of the origin, the far field fifteen away.
    mesh.label_boundary("airfoil", lambda x: np.linalg.norm(x, axis=-1) < 2.0)
    mesh.label_boundary("farfield", lambda x: np.linalg.norm(x, axis=-1) >= 2.0)
    return mesh


def make_scheme(mesh, flux):
    model = AbstractPhysicalConservativeModel(EulerPDE(dim=2, gamma=GAMMA))
    model.add_boundary_condition("airfoil", Wall())
    model.add_boundary_condition("farfield", Dirichlet(lambda _x: W_INF))
    return FiniteVolumeScheme(
        model,
        VariablesFV(mesh, nb_variables=4),
        flux,
        assemble_second_order=False,
        assemble_reaction=False,
        assemble_source=False,
    )


def wave_speed(w):
    _, velocity, _, sound = euler_primitives(w, GAMMA)
    return jnp.linalg.norm(velocity) + sound


# ── Aerodynamics from the converged state ────────────────────────────────────
def wall_pressure(mesh, dofsl):
    """Pressure of the cell behind every airfoil face, with the face midpoint,
    outward normal (into the fluid) and length.

    First order: the wall pressure is the cell's. Returned in the order of the
    faces, which is not the order along the surface -- sort by x for a plot.
    """
    faces = jnp.asarray(mesh.boundary_groups["airfoil"])
    left, right = mesh._face_fidx_to_neighbors_fidx(faces)
    inside = np.where(np.asarray(left) < 0, np.asarray(right), np.asarray(left))
    sign = np.where(np.asarray(left) < 0, -1.0, 1.0)
    weights, points = jax.vmap(mesh._local_weights_points_interface)(faces)
    normals = jax.vmap(mesh._local_unit_normals_interface)(faces)
    lengths = np.asarray(jnp.sum(weights, axis=1))
    midpoints = np.asarray(jnp.mean(points, axis=1))
    # One normal per face (straight edges); pointing out of the cell, i.e.
    # INTO the airfoil, so the force on the airfoil is + p n.
    normal = sign[:, None] * np.asarray(normals[:, 0, :])
    _, _, pressure, _ = jax.vmap(lambda w: euler_primitives(w, GAMMA))(dofsl[inside])
    return midpoints, normal, lengths, np.asarray(pressure)


def force_coefficients(midpoints, normal, lengths, pressure):
    """``C_L, C_D`` from the pressure, in the wind axes."""
    force = np.sum(((pressure - P_INF) * lengths)[:, None] * normal, axis=0)
    wind = np.array([np.cos(ALPHA), np.sin(ALPHA)])
    lift_axis = np.array([-np.sin(ALPHA), np.cos(ALPHA)])
    return force @ lift_axis / DYNAMIC_PRESSURE, force @ wind / DYNAMIC_PRESSURE


def cell_polygons(mesh):
    """Corner coordinates of every (straight-sided) cell, for a flat plot."""
    nodes, cells = np.asarray(mesh.nodes), np.asarray(mesh.cells)
    return nodes[cells[:, [0, 1, 3, 2]]]  # tensor order -> ring order


if __name__ == "__main__":
    mesh = make_mesh()
    print(
        f"NACA0012, M={MACH}, alpha={ALPHA_DEG} deg: {mesh.n_cells_total} cells, "
        f"{len(mesh.boundary_groups['airfoil'])} faces on the airfoil, "
        f"{len(mesh.boundary_groups['farfield'])} on the far field"
    )
    scheme = make_scheme(mesh, EulerHLLCFlux(GAMMA))
    initial = jnp.broadcast_to(W_INF, (mesh.n_cells_total, 4))

    print(f"pseudo-time march, HLLC, local CFL {CFL}:")
    dofsl, history = march_to_steady_state(
        scheme, initial, wave_speed, cfl=CFL, n_steps=N_STEPS, report_every=REPORT_EVERY
    )

    c_l_march, c_d_march = force_coefficients(*wall_pressure(mesh, dofsl))
    print(f"after the march: C_L = {c_l_march:.4f}   C_D = {c_d_march:.4f}")

    # ── Newton-Krylov from the marched state, on the same residual ────────
    residual_norm = jax.jit(
        lambda w: jnp.linalg.norm(FiniteVolumeScheme._assembly_scheme_pure(scheme, w))
    )
    before = float(residual_norm(dofsl))
    print(
        f"\nNewton-Krylov, {NEWTON_ITER} steps (GMRES restart {GMRES_RESTART},"
        f" {KRYLOV_CYCLES} cycles), |F| = {before:.3e}:"
    )
    solver = NewtonKrylov(
        tol=1e-6 * before,
        max_iter=NEWTON_ITER,
        cg_solver=f"gmres:{GMRES_RESTART}",
        max_iter_linear=KRYLOV_CYCLES,
    )
    start = time.time()
    solved, report = FiniteVolumeScheme.solve(
        scheme, dofsl_init=dofsl, solver=solver, return_report=True
    )
    dofsl = solved.variables.dofsl
    after = float(residual_norm(dofsl))
    print(
        f"  {int(report.n_iter)} Newton steps, {time.time() - start:.1f} s:"
        f" |F| {before:.3e} -> {after:.3e} (x {after / before:.1e})"
    )

    midpoints, normal, lengths, pressure = wall_pressure(mesh, dofsl)
    c_p = (pressure - P_INF) / DYNAMIC_PRESSURE
    c_l, c_d = force_coefficients(midpoints, normal, lengths, pressure)
    print(f"\nC_L = {c_l:.4f}   C_D = {c_d:.4f}")
    print(
        f"same fixed point as the march: dC_L = {c_l - c_l_march:+.1e},"
        f" dC_D = {c_d - c_d_march:+.1e}"
    )
    if (MACH, ALPHA_DEG) == (0.8, 1.25):
        print(
            "reference (fine grids, Vassberg-Jameson 2010): C_L = 0.3541, C_D = 0.0223"
        )
        print(
            "first order on this mesh: the shocks are there, smeared; C_D carries the"
        )
        print("numerical dissipation.")
    else:
        print("subsonic symmetric case: C_L = C_D = 0 in the continuous problem.")
    stagnation = float(c_p.max())
    isentropic = ((1 + 0.2 * MACH * MACH) ** 3.5 - 1) / (0.7 * MACH * MACH)
    print(f"C_p max = {stagnation:.3f}  (isentropic stagnation value {isentropic:.3f})")

    rho, velocity, p, sound = jax.vmap(lambda w: euler_primitives(w, GAMMA))(dofsl)
    mach = np.asarray(jnp.linalg.norm(velocity, axis=1) / sound)

    # ── Figures: Mach field near the airfoil, C_p, convergence ───────────
    # Sized to the mesh: at H_WALL = 0.003c a cell is 1.5e-3 of the field's
    # width, so the field gets a full row of a wide figure and the shock foot
    # its own zoom -- on a 5-inch panel the wall cells would be sub-pixel.
    figure = plt.figure(figsize=(18, 12), constrained_layout=True)
    grid = figure.add_gridspec(2, 3, height_ratios=(1.15, 1.0))
    ax_field = figure.add_subplot(grid[0, :])
    ax_zoom, ax_cp, ax_conv = (figure.add_subplot(grid[1, k]) for k in range(3))

    # Isolines: a Delaunay triangulation of the cell centres near the airfoil,
    # with the triangles crossing the profile masked out (their centroid is
    # inside it), so that no level line is drawn through the solid.
    centres = np.asarray(jax.vmap(mesh.cell_centroid)(jnp.arange(mesh.n_cells_total)))
    window = (np.abs(centres[:, 0] - 0.5) < 1.5) & (np.abs(centres[:, 1]) < 1.2)
    triangulation = Triangulation(centres[window, 0], centres[window, 1])
    airfoil = Path(np.vstack(naca0012_surfaces()))
    triangle_centres = centres[window][triangulation.triangles].mean(axis=1)
    triangulation.set_mask(airfoil.contains_points(triangle_centres))
    levels = np.arange(0.1, 1.4, 0.05)
    polygons = cell_polygons(mesh)

    for axis, limits, edges in (
        (ax_field, ((-0.5, 1.5), (-0.6, 0.6)), "face"),
        (ax_zoom, ((0.35, 0.85), (0.0, 0.3)), "k"),
    ):
        cells = PolyCollection(
            polygons, array=mach, cmap="jet", edgecolor=edges, lw=0.15, alpha=1.0
        )
        axis.add_collection(cells)
        contours = axis.tricontour(
            triangulation, mach[window], levels=levels, colors="k", linewidths=0.5
        )
        axis.clabel(contours, levels[::4], fontsize=7, fmt="%.1f")
        axis.tricontour(
            triangulation, mach[window], levels=[1.0], colors="w", linewidths=1.5
        )
        axis.set(xlim=limits[0], ylim=limits[1], aspect="equal")
    ax_field.set_title("Mach number, isolines every 0.05 (white: M = 1)")
    ax_zoom.set_title("the upper shock foot, cells drawn")
    figure.colorbar(cells, ax=ax_field, shrink=0.8)

    upper = midpoints[:, 1] >= 0
    for mask, label in ((upper, "upper surface"), (~upper, "lower surface")):
        order = np.argsort(midpoints[mask, 0])
        ax_cp.plot(
            midpoints[mask, 0][order], -c_p[mask][order], "-", lw=1.2, label=label
        )
    ax_cp.set(
        xlabel="x / c",
        ylabel="-C_p",
        title=f"pressure coefficient ({len(c_p)} faces), C_L={c_l:.3f}, C_D={c_d:.3f}",
    )
    ax_cp.grid(alpha=0.25)
    ax_cp.legend()

    ax_conv.semilogy(history[:, 0], history[:, 1] / history[0, 1], "o-")
    ax_conv.set(
        xlabel="pseudo-time step",
        ylabel="density residual / initial",
        title="convergence, local time step",
    )
    ax_conv.grid(alpha=0.25)
    figure.suptitle(
        f"NACA0012, M_inf = {MACH}, alpha = {ALPHA_DEG} deg, {mesh.n_cells_total} Gmsh quads"
    )
    plt.show()