# -*- coding: utf-8 -*-
# TorsoCAE Validation Journal - NACA 0012, k-omega SST, NASA TMR conditions,
# on a structured pure-hexahedron O-grid.
#
# Reference: NASA Turbulence Modeling Resource "2D NACA 0012 Airfoil
# Validation Case" -- Re_c = 6e6, M = 0.15 (incompressible here), sharp
# trailing edge closed at x = 1, far field ~500 chords, freestream
# TI = 0.052 % and nu_t/nu = 0.009. CFL3D SST on the 897x257 C-grid:
#
#     alpha   CL (CFL3D SST)   CD (CFL3D SST)
#      0        ~0              0.00809
#     10        1.0778          0.01236
#     15        1.5068          0.02219
#
# The NASA reference uses a C-grid with 513 airfoil points and an 8e-7 c
# first spacing. This journal uses NASA's scaled closed profile and 500-chord far field.
# Its O-grid topology differs from NASA's C-grid; the comparison is numerical.
#
# Why not the CSG + gmsh mesher used by cfd_naca0012_sst_validation1.py:
# gmsh's 2D boundary-layer field cannot produce these spacings here. At
# 8e-7 c it returns 50 zero-volume and 10 inverted tets (measured), which
# corrupt the heap inside the mesh partitioner instead of failing cleanly;
# 2e-6 c was the finest strictly valid unstructured mesh, and that still
# caps the far field, because a 500-chord box in a fixed-thickness slab
# means ~1e5 cell aspect ratios. A marched structured grid has neither
# limit: spacing is prescribed, and geometric stretching buys 500 chords
# for ~35 extra layers.

import math
import os
import sys
from pathlib import Path

import numpy as np

CASE_DIR = Path(__file__).resolve().parent
ROOT = CASE_DIR.parents[3]
sys.path.insert(0, str(ROOT / "server"))
from mesh_generators.structured_cgrid import build_body_in_box_mesh, estimate_grid


ALPHA_DEG = 10.0
RE = 6.0e6
CHORD = 1.0
U_INF = 1.0
RHO = 1.0
MU = RHO * U_INF * CHORD / RE

# Explicit TMR freestream turbulence keeps initialization and inlet data equal.
TURBULENCE_INTENSITY = 0.00052
VISCOSITY_RATIO = 0.009
K_INF = 1.5 * (TURBULENCE_INTENSITY * U_INF) ** 2
OMEGA_INF = RHO * K_INF / (MU * VISCOSITY_RATIO)

REFERENCE = {0.0: (0.0, 0.00809), 10.0: (1.0778, 0.01236), 15.0: (1.5068, 0.02219)}
INITIAL_STATE_NPZ = os.environ.get("TORSO_NACA_INITIAL_STATE", str(CASE_DIR / "initial_state.npz"))
# False reproduces the baseline radial O-grid with identical topology.
WAKE_REFINEMENT = True

# Surface tags the generator stamps on the mesh; the analysis case is keyed
# to these, so they double as the journal's BC map.
TAGS = {"wall": 7, "inlet": 1, "outlet": 6, "bottom": 2, "top": 4,
        "back": 3, "front": 5}

MESH = dict(
    points_around=256,      # 43 pressure nodes in the first 1% chord; 17 in the last 1%
    first_layer=8.0e-7,     # NASA first spacing; resolved viscous sublayer
    growth=1.15,            # layer-to-layer ratio inside the boundary layer
    bl_thickness=0.03,      # past this the stretching accelerates
    max_growth=1.75,        # outer growth limits RAM; BL growth remains 1.15
    # Rectangular domain (x_min, x_max, y_min, y_max) in chords, with the
    # airfoil as a hole in it.
    box=(-500.0, 500.0, -500.0, 500.0),
    thickness=0.004,        # slab span, one cell thick
    # Thins the grid without changing its structure or its wall spacing (y+
    # is unaffected): cell count falls roughly as coarsen**2. Raise this if
    # the solve does not fit in memory -- 1.0 is the full grid.
    coarsen=1.0,
)

# Direct LU: the tested iterative Schur path does not converge on this mesh.
# Actual worker RSS is about 12 GB; the grid estimator is not total RSS.
SOLVER = dict(
    algo="direct", precond="lu", tol=1e-7, max_iter=800,
    max_inner_iter=1, inner_tol=1e-3,
    n_cores=4, device="cpu",
    quadrature_degree=6,  # four Gauss points per axis for curved Q2 hexes
    live_viz_interval=10,   # complete solver fields are saved with each frame
)
# Steady NASA comparison: initialize from the freestream or an explicit saved
# field. Pseudo-time is numerical continuation, not a physical inlet ramp.
STEADY_STAGE = dict(
    steady_rans=True, time=0.0,
    pseudo_dt=10.0, pseudo_dt_max=1000.0, pseudo_dt_growth=2.0,
    pseudo_cfl=10.0, pseudo_cfl_adaptive=True, pseudo_cfl_max=50.0,
    num_steps=120, steady_tol=1e-4, steady_min_steps=20,
)
STEADY_CHUNKS = 1

# A direct (MUMPS) factorisation that does not fit is an out-of-memory kill of
# the whole machine, not a solver error, so the size is checked before the
# first solve rather than discovered during it. Raise MESH["coarsen"] if this
# trips, or raise the budget if the box really does have the memory.
MEMORY_BUDGET_GB = 6.0


def rotated_profile(alpha_deg: float, points_around: int) -> np.ndarray:
    """Closed NACA 0012 loop, nose at x = 0, rotated nose-up by alpha about
    the quarter chord (rotation is -alpha in the xy frame, so lift is +y with
    the free stream along +x). NASA profile, clustered nose, cosine tail."""
    # NASA TMR closes the original -0.1015 polynomial at 1.008930411365,
    # then scales both coordinates back to unit chord.
    t = np.linspace(0.0, 1.0, points_around // 2 + 1)
    transition = 0.8
    angle_fraction = np.where(t < transition,
        2.0 * t**2 / transition - t**3 / transition**2, t)
    x = 0.5 * (1.0 - np.cos(np.pi * angle_fraction))
    closure = 1.008930411365
    extended = closure * x
    y = 0.6 / closure * (
        0.2969 * np.sqrt(extended) - 0.1260 * extended
        - 0.3516 * extended**2 + 0.2843 * extended**3 - 0.1015 * extended**4)
    y[0] = y[-1] = 0.0
    pts = CHORD * np.vstack((np.column_stack((x[::-1], y[::-1])),
                            np.column_stack((x[1:-1], -y[1:-1]))))
    a = -math.radians(alpha_deg)
    ca, sa = math.cos(a), math.sin(a)
    xr, yr = pts[:, 0] - 0.25 * CHORD, pts[:, 1]
    return np.column_stack((0.25 * CHORD + ca * xr - sa * yr, sa * xr + ca * yr))


def redistribute_wake(mesh: dict) -> None:
    """Apply the tested NACA wake spacing, preserving topology and wall layers.

    The guide circulation is from the computed baseline CL=1.0567, not
    the NASA target. It sets mesh coordinates only, never solution values.
    All lengths below are in chord units. This case-specific
    geometry construction leaves the shared mesher and solver unchanged.
    """
    from mesh_generators.structured_cgrid import (
        _cyclic_normals, _smoothstep, box_outer_points, coarsen_grid, layer_heights,
    )

    n = int((mesh["topologies"][3]["cell_data"] == TAGS["wall"]).sum())
    n_plane = len(mesh["x"]) // 2
    base = mesh["x"][:n_plane, :2].reshape(-1, n, 2).copy() / CHORD
    profile = base[0]
    centre = profile.mean(axis=0)
    box = np.asarray(MESH["box"]) / CHORD
    outer = box_outer_points(profile, centre, box)
    _, growth, max_growth = coarsen_grid(
        profile, MESH["growth"], MESH["max_growth"], MESH["coarsen"],
    )
    heights = layer_heights(
        MESH["first_layer"] / CHORD, growth,
        float(np.linalg.norm(outer - profile, axis=1).mean()),
        max_growth, MESH["bl_thickness"] / CHORD,
    )
    if len(heights) + 1 != len(base):
        raise ValueError("Wake redistribution layer count does not match the base grid")
    normal = _cyclic_normals(profile, centre)
    angle = np.deg2rad(ALPHA_DEG)
    ca, sa = np.cos(angle), np.sin(angle)
    chord = np.array([ca, -sa])
    body = (profile - [.25, 0]) @ np.array([[ca, sa], [-sa, ca]]) + [.25, 0]
    sector = _smoothstep((body[:, 0] - .65) / (.95 - .65))
    r0, bend, spread = 20 * MESH["first_layer"] / CHORD, .003, .01
    far, out_start = float(heights[-1]), 10.0
    layers = base.copy()
    for j, radius in enumerate(np.r_[0., heights]):
        q = max(0., radius - r0)
        normal_base = bend * (-np.expm1(-q / bend))
        stream_base = q - normal_base
        normal_distance = (min(radius, r0) + normal_base
                           + spread * (np.sqrt(stream_base + bend) - np.sqrt(bend)))
        stream_distance = radius - normal_distance
        wake = profile + normal_distance * normal + stream_distance * chord
        blend = (float(_smoothstep(np.log(radius / out_start) / np.log(far / out_start)))
                 if radius > out_start else 0.)
        wake = (1 - blend) * wake + blend * (profile + (radius / far) * (outer - profile))
        layers[j] = base[j] + sector[:, None] * (wake - base[j])

    te = profile[np.argmax(body[:, 0])]
    offset = 1e-4
    xs, ys = te[0] + offset, te[1] + np.tan(-angle) * offset
    circulation = .5 * 1.0567
    g, xc = circulation / (2 * np.pi), .25
    xx = layers[:, :, 0]
    distance = np.maximum(xx - xs, 0.)
    eval_x = xs + distance
    start_radius = np.hypot(xs - xc, ys)
    psi = ys + g * np.log(start_radius)
    y = ys - g * np.log(np.hypot(eval_x - xc, ys) / start_radius)
    for _ in range(6):
        radius_squared = (eval_x - xc)**2 + y*y
        y -= (y + .5*g*np.log(radius_squared) - psi) / (1 + g*y/radius_squared)
    slope = -g * (xs - xc) / (start_radius*start_radius + g*ys)
    y += (np.tan(-angle) - slope) * distance * np.exp(-distance / .02)
    line = ys + np.tan(-angle) * distance
    shift = np.where(xx > xs, y - line, 0.)
    width = np.maximum(1., 4 * abs(shift))
    eta = (layers[:, :, 1] - line) / width
    weight = np.where(abs(eta) < 1, np.cos(np.pi*np.clip(eta, -1, 1)/2)**2, 0.)
    layers[:, :, 1] += shift * weight
    layers[0] = base[0]

    # Preserve endpoints of the box's corner-cut segments. Other outer
    # nodes slide only along their existing straight boundary face.
    gaps = np.column_stack((abs(base[-1, :, 0] - box[0]), abs(base[-1, :, 0] - box[1]),
                            abs(base[-1, :, 1] - box[2]), abs(base[-1, :, 1] - box[3])))
    side = np.argmin(gaps, axis=1)
    cuts = side != np.roll(side, -1)
    fixed = cuts | np.roll(cuts, 1)
    layers[-1, fixed] = base[-1, fixed]
    mesh["x"][:n_plane, :2] = CHORD * layers.reshape(-1, 2)
    mesh["x"][n_plane:, :2] = CHORD * layers.reshape(-1, 2)


def build_mesh() -> dict:
    m = MESH
    mesh = build_body_in_box_mesh(
        rotated_profile(ALPHA_DEG, m["points_around"]),
        box=m["box"], first_layer=m["first_layer"], growth=m["growth"],
        max_growth=m["max_growth"], bl_thickness=m["bl_thickness"],
        thickness=m["thickness"], coarsen=m["coarsen"], tags=TAGS,
    )
    if WAKE_REFINEMENT:
        redistribute_wake(mesh)
    # A structured grid can be checked outright, so check it: every corner of
    # every hexahedron must have positive volume. An inverted cell here is
    # what corrupts the mesh partitioner downstream rather than raising.
    x, hexes = mesh["x"], mesh["topologies"][5]["topology"]
    p = x[hexes]
    worst = min(
        float(np.einsum("ij,ij->i", np.cross(p[:, b] - p[:, a], p[:, c] - p[:, a]), p[:, d] - p[:, a]).min())
        for a, b, c, d in ((0, 1, 3, 4), (1, 2, 0, 5), (2, 3, 1, 6), (3, 0, 2, 7),
                           (4, 7, 5, 0), (5, 4, 6, 1), (6, 5, 7, 2), (7, 6, 4, 3))
    )
    if worst <= 0.0:
        raise RuntimeError(f"grid has inverted or degenerate hexahedra (worst corner volume {worst:.3e})")

    n_around = int((mesh["topologies"][3]["cell_data"] == TAGS["wall"]).sum())
    size = estimate_grid(n_around, len(hexes) // max(n_around, 1))
    print(f"[naca0012] grid: {len(hexes)} hexes, {len(x)} nodes, {n_around} points around, "
          f"{size['dofs']:,} dofs, ~{size['direct_solver_gb']:.1f} GB for the direct solver "
          f"(coarsen={m['coarsen']}, worst corner volume {worst:.2e})")
    if size["direct_solver_gb"] > MEMORY_BUDGET_GB:
        raise RuntimeError(
            f"this grid needs ~{size['direct_solver_gb']:.1f} GB for the direct solver, over the "
            f"{MEMORY_BUDGET_GB:.1f} GB budget. Raise MESH['coarsen'] (currently {m['coarsen']}; "
            f"cell count falls as coarsen**2, wall spacing and y+ are unaffected) or raise "
            f"MEMORY_BUDGET_GB if this machine has the memory."
        )
    return mesh


def load_state(result: dict) -> dict:
    with np.load(result["npz_path"]) as data:
        return {key: np.asarray(data[key]) for key in data.files
                if not key.startswith(("force_", "result_frame", "streamline_frame"))}


def run() -> dict:
    from pytorsocae import TorsoCAESession

    session = TorsoCAESession()
    session.set_physics("cfd", submodel="rans_k_omega_sst", backend="dolfinx")
    session.import_mesh(build_mesh(), mesh_id="mesh_0", name="NACA 0012 O-grid", body_id="body_0")
    print(f"[naca0012] boundaries: {[(b['tag'], b['elements']) for b in session.mesh_boundaries()]}")

    session.solid("NACA 0012 O-grid").material(rho=RHO, mu=MU, n_seeds=30)
    session.set_model_options(
        sst_variant="sstm",               # NASA CFL3D reference variant
        turbulence_intensity=TURBULENCE_INTENSITY,
        turbulent_viscosity_ratio=VISCOSITY_RATIO,
        wall_treatment="wall_resolved",   # y+ ~ 0.3: the sublayer is meshed, not modelled
        relax_u=0.7, relax_k=0.3, relax_omega=0.3, relax_nut=0.3,
    )
    _inflow = [U_INF, 0.0, 0.0]
    for _side in ("inlet", "bottom", "top"):
        # Velocity on the top/bottom too, not slip walls: a slip wall forbids the
        # upwash a lifting airfoil induces, which costs lift at incidence.
        session.surface(TAGS[_side], scope_id="body_0").bc(
            "inlet_velocity", values=_inflow, turbulence={"k": K_INF, "omega": OMEGA_INF},
        )
    session.surface(TAGS["outlet"], scope_id="body_0").bc("outlet_pressure", values=[0.0])
    for name in ("front", "back"):
        session.surface(TAGS[name], scope_id="body_0").bc("symmetry", components=["z"])
    session.surface(TAGS["wall"], scope_id="body_0").bc("wall")

    state = load_state({"npz_path": INITIAL_STATE_NPZ}) if INITIAL_STATE_NPZ else None
    print(f"[naca0012] initial state: {INITIAL_STATE_NPZ or 'uniform freestream'}", flush=True)
    forces, steady = [], dict(STEADY_STAGE)
    q_span = 0.5 * RHO * U_INF ** 2 * CHORD * MESH["thickness"]
    for chunk in range(STEADY_CHUNKS):
        session.set_solver_options(**SOLVER, **steady, force_tags=[TAGS["wall"]])
        result = session.compute(mesh_job_ids=["mesh_0"], initial_state=state)
        with np.load(result["npz_path"]) as data:
            forces.append(np.asarray(data["force_values"], dtype=np.float64).sum(axis=1))
        state = load_state(result)
        cl_so_far = np.concatenate(forces)[:, 1] / q_span
        print(f"[naca0012] steady chunk {chunk + 1}/{STEADY_CHUNKS}: "
              f"iterations={len(cl_so_far)} CL={cl_so_far[-1]:.4f}", flush=True)
        if result.get("cfd_steady_converged"):
            break

    forces = np.concatenate(forces, axis=0)
    cl, cd = forces[:, 1] / q_span, forces[:, 0] / q_span
    cl_final, cd_final = float(cl[-1]), float(cd[-1])
    print(f"[naca0012] iterations={len(cl)} final CL={cl_final:.4f} CD={cd_final:.5f} "
          f"(last-10 spread CL {np.ptp(cl[-10:]):.2e}, CD {np.ptp(cd[-10:]):.2e})")
    print(f"[naca0012] steady_converged={result.get('cfd_steady_converged', False)}; "
          "iteration limit alone does not establish convergence", flush=True)
    ref = REFERENCE.get(ALPHA_DEG)
    if ref is not None:
        cl_ref, cd_ref = ref
        cl_err = (cl_final - cl_ref) / cl_ref if cl_ref else cl_final
        print(f"[naca0012] reference CFL3D SST: CL={cl_ref:.4f} CD={cd_ref:.5f} -> "
              f"CL err {100 * cl_err:+.1f}% CD err {100 * (cd_final - cd_ref) / cd_ref:+.1f}%")

    session.save_project(
        f"cfd_naca0012_sst_ogrid_alpha{ALPHA_DEG:g}.tcae", output_dir=CASE_DIR / "artifacts",
        notes=(f"NACA 0012 k-omega SST on a structured pure-hex O-grid, Re {RE:.3g}, "
               f"alpha {ALPHA_DEG:g} deg: CL={cl_final:.4f} CD={cd_final:.5f} after {len(cl)} pseudo-iterations; "
               f"steady_converged={result.get('cfd_steady_converged', False)}"),
    )
    return result


if __name__ == "__main__":
    run()