"""Generic temporal-alignment helpers (metadata-INDEPENDENT).

These functions encode the project's alignment rule -- valid_date =
initialization_date + lead, aligned by ACTUAL date, never by array index -- and
the UTC-vs-local-day distinction (observed Tx is a local UTC+7 calendar-day
maximum). They use only numpy datetime64 arithmetic, so they are correct across
leap years / 29 February and carry NO assumptions about ECMWF GRIB structure.

The ECMWF-specific parts (which GRIB key holds the step, its unit, how the
control forecast is flagged) are deliberately NOT here -- they belong to the
Phase 1 preprocessing that is written only AFTER the real GRIB is inspected and
the canonical schema is frozen. This module is what that preprocessing will call
once it has parsed raw metadata into (init_datetime, lead_value, lead_unit).
"""
import numpy as np

_UNIT = {"h": "timedelta64[h]", "hours": "timedelta64[h]",
         "D": "timedelta64[D]", "d": "timedelta64[D]", "days": "timedelta64[D]",
         "m": "timedelta64[m]", "minutes": "timedelta64[m]"}


def to_datetime64(x, unit="s"):
    """Coerce scalars/arrays to datetime64 without assuming a source calendar."""
    return np.asarray(x).astype(f"datetime64[{unit}]")


def valid_time(init, lead, lead_unit="h"):
    """valid = init + lead. Aligned by real date arithmetic (handles 29 Feb).

    ``init`` datetime64 (scalar or array). ``lead`` numeric scalar/array in
    ``lead_unit`` (h/D/m). Returns datetime64 at second resolution.
    """
    if lead_unit not in _UNIT:
        raise ValueError(f"unknown lead_unit {lead_unit!r}; use one of {sorted(_UNIT)}")
    init = np.asarray(init).astype("datetime64[s]")
    dt = np.asarray(lead).astype(_UNIT[lead_unit])
    return init + dt.astype("timedelta64[s]")


def local_calendar_day(utc_times, tz_hours=7):
    """Local calendar day (datetime64[D]) each UTC timestamp belongs to.

    Observed Tx is a UTC+7 calendar-day maximum, so a forecast valid at a UTC
    instant is attributed to floor((t_utc + tz)/1 day).
    """
    t = np.asarray(utc_times).astype("datetime64[s]")
    shifted = t + np.timedelta64(int(tz_hours * 3600), "s")
    return shifted.astype("datetime64[D]")


def align_by_date(fcst_days, obs_days):
    """Map each forecast local-day to its index in ``obs_days`` by DATE, not
    position. Returns (obs_index, matched); obs_index is -1 where the date is
    absent from ``obs_days`` (e.g. a missing observation)."""
    fcst = np.asarray(fcst_days).astype("datetime64[D]")
    obs = np.asarray(obs_days).astype("datetime64[D]")
    order = np.argsort(obs)
    obs_sorted = obs[order]
    pos = np.searchsorted(obs_sorted, fcst)
    pos_clipped = np.clip(pos, 0, len(obs_sorted) - 1)
    hit = obs_sorted[pos_clipped] == fcst
    obs_index = np.where(hit, order[pos_clipped], -1)
    return obs_index, hit


def align_forecast_to_observation(init, leads, obs_days, lead_unit="h",
                                  tz_hours=7):
    """Align one forecast's leads to a local-day observation series.

    Returns a dict of equal-length arrays:
      valid_utc     - datetime64[s] valid time per lead
      local_day     - datetime64[D] local (UTC+tz) day per lead
      obs_index     - index into obs_days (-1 if that day is missing)
      matched       - bool, obs present for that local day
      partial_first - bool per lead, True for the first lead if it does not fall
                      on a local-day boundary (LT0 may be a partial day)
    """
    leads = np.asarray(leads)
    vutc = valid_time(init, leads, lead_unit)
    lday = local_calendar_day(vutc, tz_hours)
    obs_index, matched = align_by_date(lday, obs_days)
    # partial-day flag: the local-day boundary in UTC is (00 - tz) = e.g. 17:00 UTC
    secs_into_local_day = (vutc + np.timedelta64(int(tz_hours * 3600), "s")
                           - lday.astype("datetime64[s]")).astype("timedelta64[s]").astype(int)
    partial_first = np.zeros(len(leads), dtype=bool)
    if len(leads):
        partial_first[0] = secs_into_local_day[0] != 0
    return {"valid_utc": vutc, "local_day": lday, "obs_index": obs_index,
            "matched": matched, "partial_first": partial_first}