"""Weisman-Klemp analytic sounding and 1-D parcel lift utilities."""

from __future__ import annotations

import numpy as np
from scipy.integrate import cumulative_trapezoid

# --------------------------------------------------------------------------
# Physical constants
# --------------------------------------------------------------------------
G    = 9.80665        # m s⁻²
RD   = 287.04         # J kg⁻¹ K⁻¹
RV   = 461.5          # J kg⁻¹ K⁻¹
CP   = 1005.7         # J kg⁻¹ K⁻¹
LV   = 2.501e6        # J kg⁻¹
EPS  = RD / RV        # 0.6220
KAPPA = RD / CP       # 0.2854
P0   = 1000.0         # hPa reference pressure


def _sat_vapor_pressure(T_K: np.ndarray) -> np.ndarray:
    """Bolton (1980) saturation vapor pressure (hPa) from T in Kelvin."""
    T_C = T_K - 273.15
    return 6.112 * np.exp(17.67 * T_C / (T_C + 243.5))


def _dewpoint_from_qv(qv: np.ndarray, p_hPa: np.ndarray) -> np.ndarray:
    """Dewpoint (K) from mixing ratio (kg/kg) and pressure (hPa)."""
    e = p_hPa * qv / (EPS + qv)
    e = np.maximum(e, 1e-6)
    ln_term = np.log(e / 6.112)
    T_C = 243.5 * ln_term / (17.67 - ln_term)
    return T_C + 273.15


def _virtual_temp(T_K: np.ndarray, qv: np.ndarray) -> np.ndarray:
    """Virtual temperature (K)."""
    return T_K * (1.0 + qv / EPS) / (1.0 + qv)


def _sat_mixing_ratio(T_K: np.ndarray, p_hPa: np.ndarray) -> np.ndarray:
    """Saturation mixing ratio (kg/kg)."""
    es = _sat_vapor_pressure(T_K)
    return EPS * es / (p_hPa - es)


def wk_sounding(
    z: np.ndarray,
    theta_ml: float = 300.0,
    qv_ml_gkg: float = 14.0,
    z_ml_m: float = 1000.0,
    z_trop_m: float = 12000.0,
    T_trop_K: float = 213.0,
    gamma_ft: float = 1.25,
    p_sfc_hPa: float = 1000.0,
) -> dict:
    """Build the Weisman-Klemp analytic sounding on height grid ``z`` (meters AGL).

    Returns a dict with arrays: T_K, Td_K, theta, qv, rho, p_hPa (all shape z).
    """
    z = np.asarray(z, dtype=float)
    qv_ml = qv_ml_gkg * 1e-3  # convert to kg/kg

    # Estimate tropopause pressure with scale-height formula (T_mean ≈ 255 K)
    # so that θ_trop = T_trop * (P0/p_trop)^κ gives the right temperature.
    p_trop_est = p_sfc_hPa * np.exp(-G * z_trop_m / (RD * 255.0))
    theta_trop = T_trop_K * (P0 / p_trop_est) ** KAPPA

    # ---- potential temperature profile ----
    theta = np.empty_like(z)
    N2_strat = 4e-4  # N² = 4×10⁻⁴ s⁻² above tropopause
    for k, zk in enumerate(z):
        if zk <= z_ml_m:
            theta[k] = theta_ml
        elif zk <= z_trop_m:
            frac = (zk - z_ml_m) / (z_trop_m - z_ml_m)
            theta[k] = theta_ml + (theta_trop - theta_ml) * frac ** gamma_ft
        else:
            # Strong isothermal-like stratosphere: θ increases exponentially
            theta[k] = theta_trop * np.exp(N2_strat * (zk - z_trop_m) / G)

    # ---- hydrostatic pressure integration ----
    # Start from p_sfc; integrate dp/dz = -ρ g = -p g / (Rd Tv)
    # Moisture profile
    qv = np.empty_like(z)
    for k, zk in enumerate(z):
        if zk <= z_ml_m:
            qv[k] = qv_ml
        else:
            # RH = 45% in free troposphere; use preliminary T to get qvs
            # We will need to iterate once to get a consistent T and qv
            qv[k] = qv_ml  # placeholder

    # Build p(z) iteratively (one pass sufficient)
    p_hPa = np.empty_like(z)
    p_hPa[0] = p_sfc_hPa

    for k in range(1, len(z)):
        dz = z[k] - z[k - 1]
        # Mid-level T from θ and p (use previous p for first estimate)
        T_lo = theta[k - 1] * (p_hPa[k - 1] / P0) ** KAPPA
        T_hi_est = theta[k] * (p_hPa[k - 1] / P0) ** KAPPA  # first guess
        Tv_mid = _virtual_temp(0.5 * (T_lo + T_hi_est), 0.5 * (qv[k - 1] + qv[k]))
        p_hPa[k] = p_hPa[k - 1] * np.exp(-G * dz / (RD * float(Tv_mid)))

    # Recompute moisture using actual pressures (RH = 45% above BL)
    for k, zk in enumerate(z):
        T_k = theta[k] * (p_hPa[k] / P0) ** KAPPA
        if zk <= z_ml_m:
            qv[k] = qv_ml
        else:
            qvs = _sat_mixing_ratio(np.array([T_k]), np.array([p_hPa[k]]))[0]
            RH = 0.45
            qv[k] = min(RH * qvs, qv_ml)  # cap at mixed-layer value

    # Recompute p(z) with corrected moisture
    p_hPa[0] = p_sfc_hPa
    for k in range(1, len(z)):
        dz = z[k] - z[k - 1]
        T_lo = theta[k - 1] * (p_hPa[k - 1] / P0) ** KAPPA
        T_hi_est = theta[k] * (p_hPa[k - 1] / P0) ** KAPPA
        Tv_lo = _virtual_temp(np.array([T_lo]), np.array([qv[k - 1]]))[0]
        Tv_hi = _virtual_temp(np.array([T_hi_est]), np.array([qv[k]]))[0]
        Tv_mid = 0.5 * (Tv_lo + Tv_hi)
        p_hPa[k] = p_hPa[k - 1] * np.exp(-G * dz / (RD * Tv_mid))

    # Temperature and dewpoint
    T_K = theta * (p_hPa / P0) ** KAPPA
    Td_K = _dewpoint_from_qv(qv, p_hPa)
    Td_K = np.minimum(Td_K, T_K)  # Td cannot exceed T

    # Density
    Tv = _virtual_temp(T_K, qv)
    rho = p_hPa * 100.0 / (RD * Tv)  # kg m⁻³

    return {
        "T_K": T_K,
        "Td_K": Td_K,
        "theta": theta,
        "qv": qv,
        "rho": rho,
        "p_hPa": p_hPa,
    }


def _moist_lapse_rate(T_K: float, p_hPa: float) -> float:
    """Saturated adiabatic lapse rate dT/dz (K/m), negative (decreases with height)."""
    rs = float(_sat_mixing_ratio(np.array([T_K]), np.array([p_hPa]))[0])
    numer = G / CP * (1.0 + LV * rs / (RD * T_K))
    denom = 1.0 + LV ** 2 * rs / (CP * RV * T_K ** 2)
    return -numer / denom  # negative


def lift_parcel(
    z: np.ndarray,
    T_env: np.ndarray,
    qv_env: np.ndarray,
    p_hPa: np.ndarray,
    delta_T_K: float = 0.0,
) -> dict:
    """Lift a surface parcel with temperature excess ``delta_T_K`` above the environment.

    Returns dict: T_parcel, Td_parcel, LCL_m, LFC_m, EL_m, CAPE, CIN, B (buoyancy profile).
    """
    z = np.asarray(z, dtype=float)
    T_env = np.asarray(T_env, dtype=float)
    qv_env = np.asarray(qv_env, dtype=float)
    p_hPa = np.asarray(p_hPa, dtype=float)

    dz = z[1] - z[0]
    N = len(z)

    T_p0 = T_env[0] + delta_T_K
    qv_p0 = qv_env[0]  # conserve mixing ratio

    # Surface θ_e (approximately conserved during moist ascent)
    theta_p0 = T_p0 * (P0 / p_hPa[0]) ** KAPPA

    T_parcel = np.empty(N)
    T_parcel[0] = T_p0
    LCL_m = None

    # Dry adiabatic ascent until T_parcel == Td_parcel (LCL)
    above_lcl = False
    T_p = T_p0
    p_p = p_hPa[0]

    for k in range(1, N):
        if not above_lcl:
            # Dry adiabatic: conserve θ
            T_p = theta_p0 * (p_hPa[k] / P0) ** KAPPA
            # Dewpoint of parcel (mixing ratio conserved below LCL)
            Td_p = float(_dewpoint_from_qv(np.array([qv_p0]), np.array([p_hPa[k]]))[0])
            if T_p <= Td_p + 0.01 and LCL_m is None:
                LCL_m = float(z[k])
                above_lcl = True
        if above_lcl:
            # Moist adiabatic: integrate dT/dz numerically
            gamma_m = _moist_lapse_rate(T_p, float(p_hPa[k]))
            T_p = T_p + gamma_m * dz
        T_parcel[k] = T_p

    if LCL_m is None:
        LCL_m = float(z[-1])

    # Dewpoint of parcel: conserved below LCL, saturated above
    Td_parcel = np.empty(N)
    for k in range(N):
        if z[k] <= LCL_m:
            Td_parcel[k] = float(_dewpoint_from_qv(np.array([qv_p0]), np.array([p_hPa[k]]))[0])
        else:
            Td_parcel[k] = T_parcel[k]  # saturated above LCL

    # Virtual temperature correction
    qv_parcel = np.empty(N)
    for k in range(N):
        if z[k] <= LCL_m:
            qv_parcel[k] = qv_p0
        else:
            qv_parcel[k] = float(_sat_mixing_ratio(np.array([T_parcel[k]]), np.array([p_hPa[k]]))[0])

    Tv_parcel = _virtual_temp(T_parcel, qv_parcel)
    Tv_env = _virtual_temp(T_env, qv_env)

    # Buoyancy profile B(z) = g * (Tv_p - Tv_e) / Tv_e
    B = G * (Tv_parcel - Tv_env) / Tv_env

    # LFC and EL: first and last level where B > 0 above LCL
    LFC_m = LCL_m
    EL_m = LCL_m
    for k in range(N):
        if z[k] >= LCL_m and B[k] > 0:
            if LFC_m == LCL_m or z[k] < LFC_m:
                LFC_m = float(z[k])
            EL_m = float(z[k])

    # CAPE and CIN
    B_pos = np.where(B > 0, B, 0.0)
    B_neg = np.where((B < 0) & (z < LFC_m), B, 0.0)
    CAPE = float(np.trapezoid(B_pos, z))
    CIN = float(np.trapezoid(B_neg, z))

    return {
        "T_parcel": T_parcel,
        "Td_parcel": Td_parcel,
        "qv_parcel": qv_parcel,
        "LCL_m": LCL_m,
        "LFC_m": LFC_m,
        "EL_m": EL_m,
        "CAPE": CAPE,
        "CIN": CIN,
        "B": B,
    }