from __future__ import annotations
import numpy as np
import pandas as pd


def _monthly_daylight_hours(lat_deg: np.ndarray, month: int) -> np.ndarray:
    doy_mid = np.array([15,45,74,105,135,166,196,227,258,288,319,349], dtype=float)[month-1]
    phi = np.deg2rad(lat_deg)
    delta = 0.409 * np.sin(2.0 * np.pi * doy_mid / 365.0 - 1.39)
    x = np.clip(-np.tan(phi) * np.tan(delta), -1.0, 1.0)
    return 24.0 / np.pi * np.arccos(x)


def heat_index_from_monthly_clim(t2_clim: np.ndarray) -> np.ndarray:
    t = np.where(np.isfinite(t2_clim), np.maximum(t2_clim, 0.0), np.nan)
    term = (t / 5.0) ** 1.514
    return np.nansum(term, axis=0).astype(np.float32)


def thornthwaite_pet(t2: np.ndarray, lat2d: np.ndarray, time: pd.DatetimeIndex, heat_index: np.ndarray) -> np.ndarray:
    t = np.asarray(t2, dtype=float)
    I = np.asarray(heat_index, dtype=float)
    I = np.where(np.isfinite(I) & (I > 0), I, np.nan)
    a = 6.75e-7 * I**3 - 7.71e-5 * I**2 + 1.792e-2 * I + 0.49239
    out = np.full_like(t, np.nan, dtype=np.float32)
    for k, dt in enumerate(time):
        m = int(dt.month)
        days = int(dt.days_in_month)
        L = _monthly_daylight_hours(lat2d, m)
        T = t[k]
        valid = np.isfinite(T) & np.isfinite(I) & np.isfinite(a) & np.isfinite(L) & (I > 0)
        pet = np.full_like(T, np.nan, dtype=float)
        cold = valid & (T <= 0)
        warm = valid & (T > 0)
        pet[cold] = 0.0
        with np.errstate(invalid='ignore', divide='ignore', over='ignore'):
            pet[warm] = 16.0 * (L[warm]/12.0) * (days/30.0) * ((10.0*T[warm]/I[warm]) ** a[warm])
        out[k] = np.where(np.isfinite(pet) & (pet >= 0), pet, np.nan).astype(np.float32)
    return out