"""Shared feature + target construction for EZFlow v2 (OpenFOAM-RANS surrogate).

THIS IS THE SINGLE SOURCE OF TRUTH for node features and target transforms.
Both the training ETL (etl.py) and the serving path (infer_v5.py) import from
here, so a feature added in one place can never silently diverge from the other
(that duplication bug existed between the v1 data.py and inference.py).

Convention (non-dimensional, matches the CFD): U_inf=1, rho=1, L_ref=1, nu=1/Re.

Node features (NODE_FEATURE_DIM = 12), all per-node:
    [0:3]  inflow_dir            = [cos(yaw), sin(yaw), 0]
    [3]    d_wall                = nearest-surface distance / L_ref(=1)  (TRUE scale,
                                   NOT per-case-max normalized -- absolute BL scale)
    [4]    log1p(d_wall)         = near-wall resolution where BL gradients are steep
    [5]    d_plus_proxy          = log1p(clip(d_wall*sqrt(Re), 0, 3000))  Re-aware y+ surrogate
    [6:9]  nearest_wall_normal   = outward unit normal of closest wall face (0 far away)
    [9]    normal_dot_inflow     = stagnation(+1)/leeward(-1) indicator
    [10]   is_wall               = 1 if node lies on the body surface
    [11]   local_mesh_scale      = mean incident edge length (resolution awareness)

Global features (GLOBAL_DIM = 2): [log10(Re), yaw_rad]

Targets (TARGET_DIM = 7): [u, v, w, p, log1p(k), log(omega), log(nut*Re)]
    velocity/pressure are z-scored downstream; the 3 turbulence channels are put in
    log space here (heavy-tailed, span decades) and nut is made ~Re-invariant.
"""
from __future__ import annotations
import numpy as np
from scipy.spatial import cKDTree

NODE_FEATURE_DIM = 12
GLOBAL_DIM = 2
TARGET_DIM = 7
TARGET_NAMES = ["u", "v", "w", "p", "log1p_k", "log_omega", "log_nutRe"]
L_REF = 1.0
_EPS = 1e-12


def knn_edges(pos: np.ndarray, k: int = 8) -> np.ndarray:
    """Bidirectional kNN edge_index (2, E) via cKDTree (no torch_cluster dep)."""
    tree = cKDTree(pos)
    kq = min(k + 1, len(pos))
    _, nbr = tree.query(pos, k=kq)
    nbr = np.atleast_2d(nbr)
    src = np.repeat(np.arange(len(pos)), nbr.shape[1] - 1)
    dst = nbr[:, 1:].reshape(-1)
    return np.stack([np.concatenate([src, dst]), np.concatenate([dst, src])]).astype(np.int64)


def local_mesh_scale(pos: np.ndarray, edge_index: np.ndarray) -> np.ndarray:
    """Mean incident edge length per node (resolution-awareness feature)."""
    src, dst = edge_index
    d = np.linalg.norm(pos[dst] - pos[src], axis=1)
    n = pos.shape[0]
    s = np.zeros(n); c = np.zeros(n)
    np.add.at(s, dst, d); np.add.at(c, dst, 1.0)
    return s / np.maximum(c, 1.0)


def build_node_features(pos: np.ndarray, wall_pts: np.ndarray, wall_normals: np.ndarray,
                        yaw_rad: float, Re: float, edge_index: np.ndarray,
                        is_wall: np.ndarray | None = None) -> np.ndarray:
    """Construct the (N, NODE_FEATURE_DIM) node feature matrix. Geometry is encoded
    via the nearest wall point/normal (scale-consistent, L_ref=1)."""
    n = pos.shape[0]
    dx, dy = float(np.cos(yaw_rad)), float(np.sin(yaw_rad))
    inflow = np.tile([dx, dy, 0.0], (n, 1)).astype(np.float32)
    if wall_pts is not None and len(wall_pts) > 0:
        tree = cKDTree(wall_pts)
        d, idx = tree.query(pos, k=1)
        d_wall = (d / L_REF).astype(np.float32)
        nrm = wall_normals[idx].astype(np.float32) if wall_normals is not None \
            else np.zeros((n, 3), np.float32)
    else:
        d_wall = np.zeros(n, np.float32)
        nrm = np.zeros((n, 3), np.float32)
    log_d = np.log1p(d_wall)
    d_plus = np.log1p(np.clip(d_wall * np.sqrt(max(Re, 1.0)), 0.0, 3000.0)).astype(np.float32)
    ndot = (nrm @ np.array([dx, dy, 0.0], np.float32)).astype(np.float32)
    if is_wall is None:
        is_wall = (d_wall < 1e-4).astype(np.float32)
    scale = local_mesh_scale(pos, edge_index).astype(np.float32)
    x = np.concatenate([
        inflow, d_wall[:, None], log_d[:, None], d_plus[:, None],
        nrm, ndot[:, None], is_wall.astype(np.float32)[:, None], scale[:, None],
    ], axis=1).astype(np.float32)
    return x


def global_features(Re: float, yaw_rad: float) -> np.ndarray:
    return np.array([[np.log10(max(Re, 1.0)), yaw_rad]], dtype=np.float32)


def transform_targets(U: np.ndarray, p: np.ndarray, k: np.ndarray,
                      omega: np.ndarray, nut: np.ndarray, Re: float) -> np.ndarray:
    """Raw CFD fields -> (N, 7) target array (turbulence channels in log space)."""
    k = np.clip(k, 0.0, None)
    omega = np.clip(omega, _EPS, None)
    nut = np.clip(nut, 0.0, None)
    return np.concatenate([
        U.reshape(-1, 3),
        p.reshape(-1, 1),
        np.log1p(k).reshape(-1, 1),
        np.log(omega).reshape(-1, 1),
        np.log(nut * max(Re, 1.0) + _EPS).reshape(-1, 1),
    ], axis=1).astype(np.float32)


def inverse_targets(y: np.ndarray, Re: float) -> dict:
    """(N, 7) prediction -> physical fields dict {U, p, k, omega, nut}."""
    y = np.asarray(y)
    return {
        "U": y[:, 0:3],
        "p": y[:, 3],
        "k": np.expm1(y[:, 4]).clip(0.0),
        "omega": np.exp(y[:, 5]),
        "nut": (np.exp(y[:, 6]) / max(Re, 1.0)).clip(0.0),
    }