import numpy as np
import xarray as xr

from dr_gcm.precipitation import (
    aggregate_monthly_precipitation,
    NEGATIVE_PRECIPITATION_POLICY,
    clip_negative_precipitation,
    deaccumulate_daily_precipitation,
)


def test_clip_negative_precipitation_sets_all_negative_values_to_zero():
    source = xr.DataArray([-2.0, -1.0e-12, 0.0, 3.0, np.nan], dims="time", attrs={"units": "mm"})
    result = clip_negative_precipitation(source)
    np.testing.assert_allclose(result.values[:4], [0.0, 0.0, 0.0, 3.0])
    assert np.isnan(result.values[4])
    assert source.values[0] == -2.0
    assert result.attrs["negative_value_policy"] == NEGATIVE_PRECIPITATION_POLICY


def test_deaccumulate_daily_precipitation_differences_then_clips():
    cumulative = xr.DataArray(
        [0.002, 0.001, 0.004],
        dims="forecast_period",
        coords={"forecast_period": [24, 48, 72]},
        name="tp",
        attrs={"units": "m"},
    )
    result = deaccumulate_daily_precipitation(cumulative)
    np.testing.assert_allclose(result.values, [2.0, 0.0, 3.0])
    assert result.name == "tp_daily"
    assert result.attrs["units"] == "mm day-1"
    assert result.attrs["source_units"] == "m"


def test_deaccumulate_daily_precipitation_accepts_timedelta_leads():
    cumulative = xr.DataArray(
        [0.001, 0.003],
        dims="forecast_period",
        coords={"forecast_period": np.array([24, 48], dtype="timedelta64[h]")},
        attrs={"units": "m"},
    )
    np.testing.assert_allclose(
        deaccumulate_daily_precipitation(cumulative).values,
        [1.0, 2.0],
    )


def test_deaccumulate_daily_precipitation_rejects_missing_daily_step():
    cumulative = xr.DataArray(
        [0.001, 0.003],
        dims="forecast_period",
        coords={"forecast_period": [24, 72]},
        attrs={"units": "m"},
    )
    try:
        deaccumulate_daily_precipitation(cumulative)
    except ValueError as exc:
        assert "consecutive 24-hour leads" in str(exc)
    else:
        raise AssertionError("missing 48-hour step was not rejected")


def test_aggregate_monthly_precipitation_uses_interval_start_and_clips():
    valid_time = np.arange(
        np.datetime64("1981-01-02"),
        np.datetime64("1981-03-02"),
        np.timedelta64(1, "D"),
    )
    values = np.ones(valid_time.size)
    values[0] = -5.0
    daily = xr.DataArray(
        values,
        dims="forecast_period",
        coords={
            "forecast_period": np.arange(24, 24 * (valid_time.size + 1), 24),
            "valid_time": ("forecast_period", valid_time),
        },
        name="tp_daily",
        attrs={"units": "mm day-1"},
    )
    monthly = aggregate_monthly_precipitation(daily)
    np.testing.assert_allclose(monthly.values, [30.0, 28.0])
    np.testing.assert_array_equal(
        monthly.valid_time.values,
        np.array(["1981-01-01", "1981-02-01"], dtype="datetime64[ns]"),
    )
    assert monthly.attrs["units"] == "mm month-1"


def test_aggregate_monthly_precipitation_rejects_incomplete_month():
    daily = xr.DataArray(
        np.ones(2),
        dims="forecast_period",
        coords={
            "forecast_period": [24, 48],
            "valid_time": (
                "forecast_period",
                np.array(["1981-01-02", "1981-01-03"], dtype="datetime64[D]"),
            ),
        },
        attrs={"units": "mm day-1"},
    )
    try:
        aggregate_monthly_precipitation(daily)
    except ValueError as exc:
        assert "complete month" in str(exc)
    else:
        raise AssertionError("partial month was not rejected")