"""Precipitation normalization and ECMWF cumulative-field processing."""
from __future__ import annotations

import numpy as np
import xarray as xr


NEGATIVE_PRECIPITATION_POLICY = (
    "All finite precipitation values below zero are set to zero before further processing."
)


def clip_negative_precipitation(precipitation: xr.DataArray) -> xr.DataArray:
    """Return precipitation with every finite value below zero replaced by zero.

    Missing values remain missing and the input array is not modified.
    """
    result = precipitation.where(~(precipitation < 0), 0)
    result.attrs = dict(precipitation.attrs)
    result.attrs["negative_value_policy"] = NEGATIVE_PRECIPITATION_POLICY
    return result


def _lead_hours(lead: xr.DataArray) -> np.ndarray:
    if lead.ndim != 1:
        raise ValueError("Forecast lead coordinate must be one-dimensional")
    values = np.asarray(lead.values)
    if np.issubdtype(values.dtype, np.timedelta64):
        hours = values / np.timedelta64(1, "h")
    else:
        hours = values.astype(float)
    if not np.all(np.isfinite(hours)) or not np.allclose(hours, np.round(hours)):
        raise ValueError("Forecast lead values must be finite whole hours")
    return np.asarray(np.round(hours), dtype=np.int64)


def deaccumulate_daily_precipitation(
    cumulative: xr.DataArray,
    lead_dimension: str = "forecast_period",
) -> xr.DataArray:
    """Convert ECMWF cumulative total precipitation in metres to daily mm.

    The first stored value must be step 24 and represents 0--24 h. Later
    values must be consecutive 24-hour steps and are differenced from the
    preceding cumulative value. Every negative daily increment is then set
    to zero before the result is returned.
    """
    if lead_dimension not in cumulative.dims:
        raise ValueError(f"Precipitation is missing lead dimension {lead_dimension!r}")
    if lead_dimension not in cumulative.coords:
        raise ValueError(f"Precipitation is missing lead coordinate {lead_dimension!r}")
    hours = _lead_hours(cumulative[lead_dimension])
    if hours.size == 0 or hours[0] != 24 or np.any(np.diff(hours) != 24):
        raise ValueError(
            "Cumulative precipitation requires consecutive 24-hour leads beginning at step 24; "
            f"found {hours.tolist()}"
        )
    units = str(cumulative.attrs.get("units", "")).strip()
    if units != "m":
        raise ValueError(f"Cumulative precipitation units must be 'm'; found {units!r}")

    first_day = cumulative.isel({lead_dimension: slice(0, 1)})
    later_days = cumulative.diff(lead_dimension, label="upper")
    daily_mm = xr.concat([first_day, later_days], dim=lead_dimension) * 1000.0
    daily_mm = clip_negative_precipitation(daily_mm)
    daily_mm.name = f"{cumulative.name}_daily" if cumulative.name else "precipitation_daily"
    daily_mm.attrs.update(
        {
            "units": "mm day-1",
            "source_units": units,
            "temporal_processing": (
                "step 24 used for 0-24 h; subsequent days are X(t)-X(t-24 h)"
            ),
        }
    )
    return daily_mm


def aggregate_monthly_precipitation(
    daily: xr.DataArray,
    lead_dimension: str = "forecast_period",
    valid_time_coordinate: str = "valid_time",
) -> xr.DataArray:
    """Sum complete clipped daily precipitation intervals into calendar months.

    A value valid at 00 UTC represents the preceding 24-hour interval, so its
    calendar-day label is valid_time minus 24 hours. The function rejects gaps
    and incomplete months instead of silently summing partial data.
    """
    if lead_dimension not in daily.dims:
        raise ValueError(f"Daily precipitation is missing dimension {lead_dimension!r}")
    if valid_time_coordinate not in daily.coords:
        raise ValueError(
            f"Daily precipitation is missing coordinate {valid_time_coordinate!r}"
        )
    valid_time = daily[valid_time_coordinate]
    if valid_time.dims != (lead_dimension,):
        raise ValueError(
            f"{valid_time_coordinate!r} must be one-dimensional along {lead_dimension!r}"
        )
    if not np.issubdtype(valid_time.dtype, np.datetime64):
        raise ValueError(f"{valid_time_coordinate!r} must contain datetime64 values")
    units = str(daily.attrs.get("units", "")).strip()
    if units not in {"mm day-1", "mm/day"}:
        raise ValueError(
            "Daily precipitation units must be 'mm day-1' or 'mm/day'; "
            f"found {units!r}"
        )

    interval_start = (
        np.asarray(valid_time.values).astype("datetime64[ns]") - np.timedelta64(24, "h")
    )
    if interval_start.size == 0:
        raise ValueError("Daily precipitation has no valid intervals")
    if interval_start.size > 1:
        spacing_hours = np.diff(interval_start) / np.timedelta64(1, "h")
        if np.any(spacing_hours != 24):
            raise ValueError("Daily precipitation contains missing or duplicate intervals")

    month_labels = interval_start.astype("datetime64[M]")
    months, counts = np.unique(month_labels, return_counts=True)
    expected_counts = (
        (months + np.timedelta64(1, "M")).astype("datetime64[D]")
        - months.astype("datetime64[D]")
    ).astype(int)
    if np.any(counts != expected_counts):
        details = {
            str(month): {"found": int(found), "expected": int(expected)}
            for month, found, expected in zip(months, counts, expected_counts)
            if found != expected
        }
        raise ValueError(f"Daily precipitation does not cover complete month(s): {details}")

    work = clip_negative_precipitation(daily)
    work = work.assign_coords(interval_start=(lead_dimension, interval_start))
    work = work.swap_dims({lead_dimension: "interval_start"})
    for coordinate in (valid_time_coordinate, lead_dimension):
        if coordinate in work.coords:
            work = work.drop_vars(coordinate)
    monthly = work.resample(interval_start="MS").sum(skipna=False)
    monthly = monthly.rename({"interval_start": "valid_time"})
    monthly.name = daily.name.replace("_daily", "_monthly") if daily.name else "precipitation_monthly"
    monthly.attrs = dict(daily.attrs)
    monthly.attrs.update(
        {
            "units": "mm month-1",
            "negative_value_policy": NEGATIVE_PRECIPITATION_POLICY,
            "temporal_aggregation": "sum of clipped daily intervals over [month_start, next_month_start)",
        }
    )
    return monthly