"""Sounding and storm-scale diagnostic calculations."""

from __future__ import annotations

import numpy as np


def bunkers_storm_motion(
    z: np.ndarray,
    u: np.ndarray,
    v: np.ndarray,
    z_top_km: float = 6.0,
    deviation_ms: float = 7.5,
) -> dict:
    """Bunkers et al. (2000) right- and left-mover storm motion.

    Uses the 0–``z_top_km`` mean wind plus ±``deviation_ms`` m/s
    perpendicular to the 0–``z_top_km`` shear vector.

    Returns dict: rm_u, rm_v (right mover), lm_u, lm_v (left mover),
    mean_u, mean_v.
    """
    z = np.asarray(z, dtype=float)
    u = np.asarray(u, dtype=float)
    v = np.asarray(v, dtype=float)

    mask = z <= z_top_km * 1000.0
    if mask.sum() < 2:
        return {"rm_u": 0.0, "rm_v": 0.0, "lm_u": 0.0, "lm_v": 0.0,
                "mean_u": 0.0, "mean_v": 0.0}

    mean_u = float(np.trapezoid(u[mask], z[mask]) / (z[mask][-1] - z[mask][0]))
    mean_v = float(np.trapezoid(v[mask], z[mask]) / (z[mask][-1] - z[mask][0]))

    shr_u = float(u[mask][-1] - u[mask][0])
    shr_v = float(v[mask][-1] - v[mask][0])
    mag = max(np.hypot(shr_u, shr_v), 1e-6)

    # Right mover: + deviation perpendicular to shear (rotated +90°)
    rm_u = mean_u + deviation_ms * shr_v / mag
    rm_v = mean_v - deviation_ms * shr_u / mag
    # Left mover: − deviation
    lm_u = mean_u - deviation_ms * shr_v / mag
    lm_v = mean_v + deviation_ms * shr_u / mag

    return {
        "rm_u": rm_u, "rm_v": rm_v,
        "lm_u": lm_u, "lm_v": lm_v,
        "mean_u": mean_u, "mean_v": mean_v,
    }


def srh(
    z: np.ndarray,
    u: np.ndarray,
    v: np.ndarray,
    storm_u: float,
    storm_v: float,
    z_bot_m: float,
    z_top_m: float,
) -> float:
    """Storm-Relative Helicity (m² s⁻²) for the specified layer.

    Uses the discrete hodograph-area formula:
        SRH = Σ [(u_{n+1} − cu)(v_n − cv) − (u_n − cu)(v_{n+1} − cv)]
    """
    z = np.asarray(z, dtype=float)
    u = np.asarray(u, dtype=float)
    v = np.asarray(v, dtype=float)

    mask = (z >= z_bot_m) & (z <= z_top_m)
    u_l = u[mask] - storm_u
    v_l = v[mask] - storm_v

    if len(u_l) < 2:
        return 0.0

    return float(np.sum((u_l[1:] - u_l[:-1]) * (v_l[1:] + v_l[:-1]) -
                        (v_l[1:] - v_l[:-1]) * (u_l[1:] + u_l[:-1])) * 0.5)


def updraft_helicity(
    w3d: np.ndarray,
    zeta3d: np.ndarray,
    z: np.ndarray,
    z_bot_m: float,
    z_top_m: float,
) -> float:
    """Updraft Helicity UH = ∫ w ζ_z dz over the updraft core center column."""
    z = np.asarray(z, dtype=float)
    mask = (z >= z_bot_m) & (z <= z_top_m)
    if mask.sum() < 2:
        return 0.0

    cx = w3d.shape[0] // 2
    cy = w3d.shape[1] // 2
    w_col = w3d[cx, cy, mask]
    z_col = zeta3d[cx, cy, mask]
    return float(np.trapezoid(w_col * z_col, z[mask]))


def bulk_wind_shear(z: np.ndarray, u: np.ndarray, v: np.ndarray,
                   z_bot_m: float, z_top_m: float) -> float:
    """Bulk wind shear magnitude (m/s) in a layer."""
    z = np.asarray(z, dtype=float)
    u = np.asarray(u, dtype=float)
    v = np.asarray(v, dtype=float)
    u_bot = float(np.interp(z_bot_m, z, u))
    v_bot = float(np.interp(z_bot_m, z, v))
    u_top = float(np.interp(z_top_m, z, u))
    v_top = float(np.interp(z_top_m, z, v))
    return float(np.hypot(u_top - u_bot, v_top - v_bot))


def collect_diagnostics(
    z: np.ndarray,
    snd: dict,
    parcel: dict,
    u_hodo: np.ndarray,
    v_hodo: np.ndarray,
    z_hodo: np.ndarray,
    w3d: np.ndarray,
    zeta3d: np.ndarray,
) -> dict:
    """Compute and return all sounding/storm-scale diagnostics as a flat dict."""
    z_hodo = np.asarray(z_hodo, dtype=float) * 1000.0  # km → m
    u_env = np.interp(z, z_hodo, np.asarray(u_hodo, dtype=float))
    v_env = np.interp(z, z_hodo, np.asarray(v_hodo, dtype=float))

    bunk = bunkers_storm_motion(z_hodo, u_hodo, v_hodo)

    srh_02 = srh(z_hodo, u_hodo, v_hodo, bunk["rm_u"], bunk["rm_v"], 0.0, 2000.0)
    srh_25 = srh(z_hodo, u_hodo, v_hodo, bunk["rm_u"], bunk["rm_v"], 2000.0, 5000.0)
    uh_02  = updraft_helicity(w3d, zeta3d, z, 0.0, 2000.0)
    uh_25  = updraft_helicity(w3d, zeta3d, z, 2000.0, 5000.0)
    bws_06 = bulk_wind_shear(z_hodo, u_hodo, v_hodo, 0.0, 6000.0)

    cx = w3d.shape[0] // 2
    cy = w3d.shape[1] // 2
    w_max = float(np.max(w3d[cx, cy, :]))

    return {
        "CAPE": parcel["CAPE"],
        "CIN": parcel["CIN"],
        "LCL_m": parcel["LCL_m"],
        "LFC_m": parcel["LFC_m"],
        "EL_m": parcel["EL_m"],
        "z_top_m": parcel.get("z_top_m", parcel["EL_m"]),
        "SRH_02": srh_02,
        "SRH_25": srh_25,
        "UH_02": uh_02,
        "UH_25": uh_25,
        "BWS_06": bws_06,
        "w_max": w_max,
        "storm_u": bunk["rm_u"],
        "storm_v": bunk["rm_v"],
        "lm_u": bunk["lm_u"],
        "lm_v": bunk["lm_v"],
        "mean_u": bunk["mean_u"],
        "mean_v": bunk["mean_v"],
    }