"""
Quantities CALCULATED from 3A sensor data, one implementation for the whole
pipeline. Every block imports them from here and never keeps its own formula,
so a value (e.g. the wet-bulb temperature of an hour) is the same in every
block and in every report.

Reference [3A-DOC]: 3A, "Algoritmi di stima di alcune grandezze
agrometeorologiche" (docs/ImplementazioneAlgoritmiSensori.pdf), section 4
"Temperatura bulbo umido (Tw)". It is the formula 3A uses for its own TBulbMed:
on (hourly mean temperature, mean humidity) it gives 3A's TBulbMed in 49 of 49
hours at station 566 (to 3A's 0.1 °C rounding, checked 2026-10-21).

Wet-bulb of one hour, from the 3A hourly columns (add_wet_bulb):
    WetBulbMin = Tw(TemperaturaAriaMin, UmiditaRelativaMax)   lowest of the hour
    WetBulbMax = Tw(TemperaturaAriaMax, UmiditaRelativaMin)   highest of the hour
Why these pairs: within an hour the water vapour in the air changes little, so
the coldest moment is also the most humid one (and the warmest the driest).
The same pairing (Tmin with RHmax, Tmax with RHmin) is used by the 3A document
for the vapour pressure, section 3 eq. (2) (FAO-56). Checked on 566: the vapour
pressure of (Tmin, RHmax), (Tmean, RHmean) and (Tmax, RHmin) differs by < 1 hPa.

Missing values: no hourly minimum / maximum temperature -> hourly mean
temperature; no hourly maximum / minimum humidity -> hourly mean humidity; no
temperature or no humidity at all -> NaN (the block decides what to do).

    from damage_engine.sources import derived
    df = fetch_hourly(ids, start, end, [*derived.WET_BULB_INPUTS, ...])
    df = derived.add_wet_bulb(df)          # adds WetBulbMin, WetBulbMax (°C)
"""

import numpy as np
import pandas as pd

K0 = 273.15                 # 0 °C in K

# constants of the 3A document, section 4 [3A-DOC]
RV = 0.11                   # cal/(Kg K)  gas constant of moist air
CP = 0.24                   # cal/(Kg K)  specific heat at constant pressure
L_EVAPORATION = 597.0       # cal/Kg      latent heat, T > 0 °C (water)
L_SUBLIMATION = 677.0       # cal/Kg      latent heat, T <= 0 °C (ice)
PRESSURE_HPA = 1000.0       # mbar        atmospheric pressure ("generalmente si può porre 1000 mbar")
# saturation vapour pressure e*(T) = 6.1 * 10^(a*T / (T + b)), T in °C
MAGNUS_WATER = (7.5, 237.3)     # T > 0 °C
MAGNUS_ICE = (9.5, 265.5)       # T <= 0 °C

# 3A hourly columns (api_parameters) used, and the columns added by add_wet_bulb
T_MIN, T_MEAN, T_MAX = "TemperaturaAriaMin", "TemperaturaAriaMedia", "TemperaturaAriaMax"
RH_MIN, RH_MEAN, RH_MAX = "UmiditaRelativaMin", "UmiditaRelativaMedia", "UmiditaRelativaMax"
WET_BULB_INPUTS = [T_MIN, T_MEAN, T_MAX, RH_MIN, RH_MEAN, RH_MAX]
WET_BULB_MIN = "WetBulbMin"     # °C, lowest wet-bulb of the hour
WET_BULB_MAX = "WetBulbMax"     # °C, highest wet-bulb of the hour


def _branch(t_c):
    """(a, b, L) per value: ice constants for T <= 0 °C, water constants above."""
    ice = t_c <= 0
    a = np.where(ice, MAGNUS_ICE[0], MAGNUS_WATER[0])
    b = np.where(ice, MAGNUS_ICE[1], MAGNUS_WATER[1])
    latent = np.where(ice, L_SUBLIMATION, L_EVAPORATION)
    return a, b, latent


def vapour_pressure_hpa(temp_c, rh_pct):
    """Vapour pressure e = RH * e*(T) / 100 (hPa) [3A-DOC sec. 4]."""
    t = np.asarray(temp_c, dtype=float)
    rh = np.asarray(rh_pct, dtype=float)
    a, b, _ = _branch(t)
    return rh / 100.0 * 6.1 * 10.0 ** (a * t / (t + b))


def dew_point_c(temp_c, rh_pct):
    """Dew point (frost point for T <= 0 °C), °C [3A-DOC sec. 4 eq. 2-3, sec. 5]."""
    t = np.asarray(temp_c, dtype=float)
    a, b, _ = _branch(t)
    with np.errstate(divide="ignore", invalid="ignore"):
        x = np.log10(vapour_pressure_hpa(t, rh_pct) / 6.1)
        return b * x / (a - x)


def wet_bulb_c(temp_c, rh_pct, pressure_hpa=PRESSURE_HPA):
    """Wet-bulb temperature (°C) [3A-DOC sec. 4 eq. 1]:

        Tw = Td + (T - Td) / (1 + L² s / (Rv Cp T²)) - 273.15     (T, Td in K)
        s  = 0.622 e / (P - 0.378 e)                              (specific humidity)
    Water constants for T > 0 °C, ice constants for T <= 0 °C. RH <= 0 gives NaN.
    """
    t = np.asarray(temp_c, dtype=float)
    _, _, latent = _branch(t)
    e = vapour_pressure_hpa(t, rh_pct)
    with np.errstate(divide="ignore", invalid="ignore"):
        s = 0.622 * e / (pressure_hpa - 0.378 * e)
        t_k, td_k = t + K0, dew_point_c(t, rh_pct) + K0
        return td_k + (t_k - td_k) / (1.0 + latent ** 2 * s / (RV * CP * t_k ** 2)) - K0


def _col(df, name):
    return pd.to_numeric(df[name], errors="coerce") if name in df else pd.Series(np.nan, index=df.index)


def add_wet_bulb(df: pd.DataFrame) -> pd.DataFrame:
    """Copy of the hourly rows with WetBulbMin and WetBulbMax (°C) added (see module doc)."""
    df = df.copy()
    t_mean, rh_mean = _col(df, T_MEAN), _col(df, RH_MEAN)
    t_lo, rh_hi = _col(df, T_MIN).fillna(t_mean), _col(df, RH_MAX).fillna(rh_mean)
    t_hi, rh_lo = _col(df, T_MAX).fillna(t_mean), _col(df, RH_MIN).fillna(rh_mean)
    df[WET_BULB_MIN] = np.round(wet_bulb_c(t_lo, rh_hi), 2)
    df[WET_BULB_MAX] = np.round(wet_bulb_c(t_hi, rh_lo), 2)
    return df