src/models/farquhar_model.py

"""
FarquharModel: mechanistic photosynthesis (Farquhar et al. 1980) with
Greer & Weedon (2012) temperature response for Vitis vinifera cv. Semillon.
Uses only on-site sensor inputs (PAR, Tleaf, CO2, VPD, Tair, etc.).

Parameters calibrated from:
  Greer, D.H. & Weedon, M.M. (2012) Modelling photosynthetic responses to
  temperature of grapevine (Vitis vinifera cv. Semillon) leaves on vines grown
  in a hot climate. Plant, Cell & Environment 35, 1050-1064.
"""

from typing import Optional

import numpy as np
import pandas as pd

# Gas constant J/(mol·K)
R = 8.314
# O2 concentration at chloroplast (mmol mol-1)
OI = 210.0
# Curvature of J vs light (dimensionless)
THETA = 0.9
# Quantum efficiency (mol e- per mol photons) base; PRI can scale this
ALPHA_DEFAULT = 0.24
# Dark respiration as fraction of Vcmax
RD_FRAC = 0.015

# --- Greer & Weedon (2012) Table / Fig 11 fitted parameters ---
# Cc-based values from paper: k25_Vcmax=38.5, k25_Jmax=98.3
# Ha/Hd (activation/deactivation energies) and Topt from Arrhenius fitting.
# NOTE: We use Ci-based apparent k25 values (60.0, 120.0) scaled ~1.5x from the
# paper's Cc-based values to compensate for mesophyll conductance (gm=5-10
# µmol/m²/s/Pa, paper p.1054) not modelled explicitly. The temperature SHAPE
# (Ha, Hd, Topt) is preserved from the paper.
_GW12_VCMAX = dict(k25=60.0, Ha=87700.0, Hd=203500.0, Topt=312.15)
_GW12_JMAX = dict(k25=120.0, Ha=63500.0, Hd=202900.0, Topt=309.05)


def _entropy_from_topt(Ha: float, Hd: float, Topt: float) -> float:
    """Derive entropy term S from Topt using Medlyn et al. (2002) Eqn 9."""
    return (Hd / Topt) + R * np.log(Ha / (Hd - Ha))


# Pre-compute entropy terms
_GW12_VCMAX["S"] = _entropy_from_topt(
    _GW12_VCMAX["Ha"], _GW12_VCMAX["Hd"], _GW12_VCMAX["Topt"]
)
_GW12_JMAX["S"] = _entropy_from_topt(
    _GW12_JMAX["Ha"], _GW12_JMAX["Hd"], _GW12_JMAX["Topt"]
)


def _modified_arrhenius(Tk: float, k25: float, Ha: float, Hd: float, S: float) -> float:
    """Modified Arrhenius with high-temperature deactivation (Medlyn et al. 2002 Eqn 8).

    k(T) = k25 * exp(Ha*(Tk-298.15)/(298.15*R*Tk))
                * (1 + exp((298.15*S - Hd)/(298.15*R)))
                / (1 + exp((S*Tk - Hd)/(R*Tk)))
    """
    # Activation component
    exp_ha = np.exp(Ha * (Tk - 298.15) / (298.15 * R * Tk))
    # Deactivation at reference temperature (normalisation)
    denom_25 = 1.0 + np.exp((298.15 * S - Hd) / (298.15 * R))
    # Deactivation at leaf temperature
    denom_tk = 1.0 + np.exp((S * Tk - Hd) / (R * Tk))
    return k25 * exp_ha * denom_25 / denom_tk


class FarquharModel:
    """
    Farquhar et al. (1980) with Bernacchi et al. (2001) temperature functions
    and Greer & Weedon (2012) Vcmax/Jmax for grapevine.
    """

    def __init__(
        self,
        k25_vcmax: float = _GW12_VCMAX["k25"],
        k25_jmax: float = _GW12_JMAX["k25"],
        Ha_vcmax: float = _GW12_VCMAX["Ha"],
        Hd_vcmax: float = _GW12_VCMAX["Hd"],
        S_vcmax: float = _GW12_VCMAX["S"],
        Ha_jmax: float = _GW12_JMAX["Ha"],
        Hd_jmax: float = _GW12_JMAX["Hd"],
        S_jmax: float = _GW12_JMAX["S"],
        alpha: float = ALPHA_DEFAULT,
        theta: float = THETA,
        rd_frac: float = RD_FRAC,
    ):
        self.params = {
            "k25_vcmax": k25_vcmax,
            "k25_jmax": k25_jmax,
            "Ha_vcmax": Ha_vcmax,
            "Hd_vcmax": Hd_vcmax,
            "S_vcmax": S_vcmax,
            "Ha_jmax": Ha_jmax,
            "Hd_jmax": Hd_jmax,
            "S_jmax": S_jmax,
            "alpha": alpha,
            "theta": theta,
            "rd_frac": rd_frac,
        }

    @staticmethod
    def calc_Kc(Tk: float) -> float:
        """Michaelis constant for CO2 (Bernacchi et al. 2001), in ppm scale."""
        return np.exp(38.05 - 79430.0 / (R * Tk))

    @staticmethod
    def calc_Ko(Tk: float) -> float:
        """Michaelis constant for O2 (Bernacchi et al. 2001)."""
        return np.exp(20.30 - 36380.0 / (R * Tk)) * 1000.0  # scale to match OI

    @staticmethod
    def calc_gamma_star(Tk: float) -> float:
        """CO2 compensation point (Bernacchi et al. 2001), ppm."""
        return np.exp(19.02 - 37830.0 / (R * Tk))

    def calc_Vcmax(self, Tk: float) -> float:
        """Vcmax at leaf temperature (modified Arrhenius, Greer & Weedon 2012)."""
        return _modified_arrhenius(
            Tk,
            self.params["k25_vcmax"],
            self.params["Ha_vcmax"],
            self.params["Hd_vcmax"],
            self.params["S_vcmax"],
        )

    def calc_Jmax(self, Tk: float) -> float:
        """Jmax at leaf temperature (modified Arrhenius, Greer & Weedon 2012)."""
        return _modified_arrhenius(
            Tk,
            self.params["k25_jmax"],
            self.params["Ha_jmax"],
            self.params["Hd_jmax"],
            self.params["S_jmax"],
        )

    def calc_electron_transport(self, PAR: float, Jmax: float) -> float:
        """Solve theta*J^2 - (alpha*PFD + Jmax)*J + alpha*PFD*Jmax = 0 for J."""
        alpha = self.params["alpha"]
        theta = self.params["theta"]
        if PAR <= 0 or Jmax <= 0:
            return 0.0
        pfd = PAR  # umol photons m-2 s-1
        b = alpha * pfd + Jmax
        c = alpha * pfd * Jmax
        disc = b * b - 4 * theta * c
        if disc < 0:
            return min(alpha * pfd, Jmax)
        j = (b - np.sqrt(disc)) / (2 * theta)
        return float(np.clip(j, 0, Jmax))

    def calc_CWSI(
        self,
        Tleaf: float,
        Tair: float,
        VPD: float,
        dTmin: Optional[float] = None,
        dTmax: Optional[float] = None,
    ) -> float:
        """Crop Water Stress Index. dTmin/dTmax from data or defaults."""
        dT = Tleaf - Tair
        if dTmin is None:
            dTmin = -2.0
        if dTmax is None:
            dTmax = 8.0
        if dTmax <= dTmin:
            return 0.0
        cwsi = (dT - dTmin) / (dTmax - dTmin)
        return float(np.clip(cwsi, 0.0, 1.0))

    def _ci_from_ca(self, ca: float, VPD: float, CWSI: float = 0.0) -> float:
        """Intercellular CO2 from ambient; gs reduced by VPD and CWSI.

        Calibrated so ci/ca ~ 0.7 at low VPD (Greer & Weedon 2012 Fig 2c),
        declining with increasing VPD and water stress.
        """
        vpd_scale = np.exp(-0.3 * max(0, VPD - 1.0)) if VPD is not None else 1.0
        stress = 1.0 - 0.5 * (CWSI if CWSI is not None else 0.0)
        gs_factor = 2.1 * vpd_scale * stress
        if gs_factor <= 0:
            return ca * 0.3
        ci = ca * (1.0 - 1.0 / (1.6 * gs_factor))
        return float(np.clip(ci, ca * 0.3, ca))

    def _compute_rates(
        self,
        PAR: float,
        Tleaf: float,
        CO2: float,
        VPD: float,
        CWSI: Optional[float] = None,
    ) -> tuple[float, float, float]:
        """Shared FvCB core: compute Rubisco-limited (Ac), RuBP-limited (Aj), and dark respiration (Rd).

        Returns (Ac, Aj, Rd) — all in umol CO2 m-2 s-1.
        """
        Tk = Tleaf + 273.15
        Kc = self.calc_Kc(Tk)
        Ko = self.calc_Ko(Tk)
        gamma_star = self.calc_gamma_star(Tk)
        Vcmax = self.calc_Vcmax(Tk)
        Jmax = self.calc_Jmax(Tk)
        J = self.calc_electron_transport(PAR, Jmax)
        Rd = self.params["rd_frac"] * Vcmax
        ci = self._ci_from_ca(CO2, VPD, CWSI)

        Ac = Vcmax * (ci - gamma_star) / (ci + Kc * (1.0 + OI / Ko))
        Aj = J * (ci - gamma_star) / (4.0 * ci + 8.0 * gamma_star)
        return Ac, Aj, Rd

    def calc_photosynthesis(
        self,
        PAR: float,
        Tleaf: float,
        CO2: float,
        VPD: float,
        Tair: float,
        CWSI: Optional[float] = None,
    ) -> float:
        """
        Net assimilation A (umol CO2 m-2 s-1). PAR in umol m-2 s-1, T in degC,
        CO2 in ppm, VPD in kPa.
        """
        Ac, Aj, Rd = self._compute_rates(PAR, Tleaf, CO2, VPD, CWSI)
        A = min(Ac, Aj) - Rd
        return float(max(0.0, A))

    def calc_photosynthesis_semillon(
        self,
        PAR: float,
        Tleaf: float,
        CO2: float,
        VPD: float,
        Tair: float,
        CWSI: Optional[float] = None,
        transition_temp: Optional[float] = None,
    ) -> tuple[float, str, bool]:
        """
        FvCB with explicit Semillon state transition.

        Returns (A, limiting_state, shading_helps):
          - A: net assimilation (umol CO2 m-2 s-1), clipped to >= 0
          - limiting_state: "RuBP_Limited" or "Rubisco_Limited"
          - shading_helps: True ONLY when the vine is Rubisco-limited AND
            light is abundant relative to Vcmax capacity (Aj > Ac), meaning
            reducing PAR would lower Aj toward Ac without reducing A.
            When False, shading would reduce A — keep panels tracking.

        Parameters
        ----------
        PAR : float
            Photosynthetically active radiation (umol photons m-2 s-1).
        Tleaf : float
            Leaf temperature (°C).
        CO2 : float
            Ambient CO2 (ppm).
        VPD : float
            Vapour pressure deficit (kPa).
        Tair : float
            Air temperature (°C).
        CWSI : float, optional
            Crop Water Stress Index (0–1). Default 0 (no stress).
        transition_temp : float, optional
            Semillon RuBP→Rubisco transition temperature (°C).
            Default from config: SEMILLON_TRANSITION_TEMP_C.
        """
        if transition_temp is None:
            from config.settings import SEMILLON_TRANSITION_TEMP_C
            transition_temp = SEMILLON_TRANSITION_TEMP_C

        Ac, Aj, Rd = self._compute_rates(PAR, Tleaf, CO2, VPD, CWSI or 0.0)
        An = min(Ac, Aj) - Rd

        # Smooth sigmoid transition (28–32°C) instead of hard threshold.
        # rubisco_weight = 0 below 28°C, 1 above 32°C, sigmoid in between.
        import math
        _TRANSITION_WIDTH = 2.0  # °C half-width of sigmoid zone
        rubisco_weight = 1.0 / (1.0 + math.exp(-(Tleaf - transition_temp) / (_TRANSITION_WIDTH / 2.5)))

        if rubisco_weight < 0.3:
            state = "RuBP_Limited"
        elif rubisco_weight > 0.7:
            state = "Rubisco_Limited"
        else:
            state = "Transition"

        # shading_helps is weighted: only meaningful when Rubisco-limited
        # AND light is abundant relative to enzyme capacity (Aj > Ac)
        shading_helps = rubisco_weight > 0.5 and (Aj > Ac)

        return float(max(0.0, An)), state, shading_helps

    def compute_all(
        self,
        df: pd.DataFrame,
        par_col: str = "Air1_PAR_ref",
        tleaf_col: str = "Air1_leafTemperature_ref",
        co2_col: str = "Air1_CO2_ref",
        vpd_col: str = "Air1_VPD_ref",
        tair_col: str = "Air1_airTemperature_ref",
        humidity_col: Optional[str] = "Air1_airHumidity_ref",
    ) -> pd.Series:
        """
        Compute A for each row using vectorized pandas operations (~100x faster).
        Returns Series of A (umol CO2 m-2 s-1), index aligned to df.
        """
        required = [par_col, tleaf_col, co2_col, vpd_col, tair_col]
        for c in required:
            if c not in df.columns:
                return pd.Series(np.nan, index=df.index)

        # Extract columns as float arrays
        par = df[par_col].astype(float)
        tleaf = df[tleaf_col].astype(float)
        co2 = df[co2_col].astype(float)
        vpd = df[vpd_col].astype(float)
        tair = df[tair_col].astype(float)

        # Vectorized CWSI from (Tleaf - Tair) with empirical bounds
        dT = tleaf - tair
        n_valid = dT.notna().sum()
        dTmin = float(dT.quantile(0.05)) if n_valid > 10 else -2.0
        dTmax = float(dT.quantile(0.95)) if n_valid > 10 else 8.0
        if dTmax <= dTmin:
            cwsi = pd.Series(0.0, index=df.index)
        else:
            cwsi = ((dT - dTmin) / (dTmax - dTmin)).clip(0.0, 1.0)

        # Vectorized FvCB computation
        Tk = tleaf + 273.15

        # Michaelis constants (Bernacchi et al. 2001) — vectorized
        Kc = np.exp(38.05 - 79430.0 / (R * Tk))
        Ko = np.exp(20.30 - 36380.0 / (R * Tk)) * 1000.0
        gamma_star = np.exp(19.02 - 37830.0 / (R * Tk))

        # Vcmax and Jmax (modified Arrhenius) — vectorized
        Vcmax = _modified_arrhenius(Tk, self.params["k25_vcmax"], self.params["Ha_vcmax"],
                                     self.params["Hd_vcmax"], self.params["S_vcmax"])
        Jmax = _modified_arrhenius(Tk, self.params["k25_jmax"], self.params["Ha_jmax"],
                                    self.params["Hd_jmax"], self.params["S_jmax"])

        # Electron transport J — vectorized quadratic solve
        alpha = self.params["alpha"]
        theta = self.params["theta"]
        b = alpha * par + Jmax
        c_val = alpha * par * Jmax
        disc = b * b - 4 * theta * c_val
        disc_safe = disc.clip(lower=0)
        J = ((b - np.sqrt(disc_safe)) / (2 * theta)).clip(lower=0)
        J = np.minimum(J, Jmax)
        J = J.where(par > 0, 0.0)

        # Dark respiration
        Rd = self.params["rd_frac"] * Vcmax

        # Intercellular CO2 — vectorized ci/ca
        vpd_scale = np.exp(-0.3 * (vpd - 1.0).clip(lower=0))
        stress = 1.0 - 0.5 * cwsi.fillna(0)
        gs_factor = 2.1 * vpd_scale * stress
        ci = co2 * (1.0 - 1.0 / (1.6 * gs_factor.clip(lower=0.01)))
        ci = ci.clip(lower=co2 * 0.3, upper=co2)

        # Rubisco-limited (Ac) and RuBP-limited (Aj) rates
        Ac = Vcmax * (ci - gamma_star) / (ci + Kc * (1.0 + OI / Ko))
        Aj = J * (ci - gamma_star) / (4.0 * ci + 8.0 * gamma_star)

        # Net assimilation
        An = np.minimum(Ac, Aj) - Rd
        An = An.clip(lower=0.0)

        # NaN where any input was NaN
        valid = par.notna() & tleaf.notna() & co2.notna() & vpd.notna() & tair.notna()
        An = An.where(valid, np.nan)

        return An