from __future__ import annotations

from typing import Iterable

import numpy as np

from src.safe_ops import safe_mean, valid_temperature_c


def _monthly_daylight_hours(lat_deg: np.ndarray, month: int) -> np.ndarray:
    """Approximate mean daylight hours for the 15th day of a calendar month."""
    day_of_year_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 * day_of_year_mid / 365.0 - 1.39)
    x = -np.tan(phi) * np.tan(delta)
    x = np.clip(x, -1.0, 1.0)
    omega_s = np.arccos(x)
    return 24.0 / np.pi * omega_s


def thornthwaite_heat_index_from_climatology(t2_clim_c: np.ndarray, min_valid_months: int = 6) -> np.ndarray:
    """Compute Thornthwaite annual heat index I from monthly climatological T2.

    Missing values and physically unreasonable temperatures are not allowed to
    become valid zero-temperature months.  Months with T <= 0 degC contribute
    zero heat, but invalid/missing months are excluded.  If fewer than
    min_valid_months are valid, I is missing.
    """
    t = valid_temperature_c(t2_clim_c)
    valid_count = np.sum(np.isfinite(t), axis=0)
    temp_pos = np.where(np.isfinite(t), np.maximum(t, 0.0), np.nan)
    term = np.where(np.isfinite(temp_pos), (temp_pos / 5.0) ** 1.514, np.nan)
    I = np.sum(np.where(np.isfinite(term), term, 0.0), axis=0)
    I = np.where(valid_count >= int(min_valid_months), I, np.nan)
    I = np.where(np.isfinite(I), I, np.nan)
    return I.astype(np.float32)


def thornthwaite_pet_monthly(
    t2_c: np.ndarray,
    lat2d: np.ndarray,
    months: Iterable[int],
    days_in_month: Iterable[int],
    heat_index: np.ndarray,
) -> np.ndarray:
    """Compute monthly Thornthwaite PET in mm/month.

    The function distinguishes physical cold months from missing/invalid data:
    * valid T <= 0 degC => PET = 0 by Thornthwaite convention;
    * missing/invalid T, invalid heat index, invalid latitude/day count => PET = NaN.
    """
    t2 = valid_temperature_c(t2_c)
    months = np.asarray(list(months), dtype=int)
    ndays = np.asarray(list(days_in_month), dtype=float)
    if t2.shape[0] != months.size or months.size != ndays.size:
        raise ValueError('time length mismatch among t2, months, and days_in_month')
    I = np.asarray(heat_index, dtype=float)
    I = np.where(np.isfinite(I) & (I > 0.0), I, np.nan)
    lat2d = np.asarray(lat2d, dtype=float)
    a = 6.75e-7 * I**3 - 7.71e-5 * I**2 + 1.792e-2 * I + 0.49239
    out = np.full_like(t2, np.nan, dtype=np.float32)
    for k, m in enumerate(months):
        if not (1 <= int(m) <= 12) or not np.isfinite(ndays[k]) or ndays[k] <= 0:
            continue
        T = t2[k]
        L = _monthly_daylight_hours(lat2d, int(m))
        valid_base = 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_base & (T <= 0.0)
        pet[cold] = 0.0
        warm = valid_base & (T > 0.0)
        with np.errstate(invalid='ignore', divide='ignore', over='ignore'):
            pet[warm] = 16.0 * (L[warm] / 12.0) * (ndays[k] / 30.0) * ((10.0 * T[warm] / I[warm]) ** a[warm])
        pet[~np.isfinite(pet)] = np.nan
        pet = np.where(pet >= 0.0, pet, np.nan)
        out[k] = pet.astype(np.float32)
    return out