import numpy as np
import xarray as xr

from dr_gcm.monthly import build_monthly_dataset, instantaneous_monthly, run


INIT = np.datetime64("1981-01-01T00:00:00", "ns")
BOUNDARY = np.datetime64("1981-08-01T00:00:00", "ns")


def instant_dataset():
    valid = np.arange(INIT + np.timedelta64(6, "h"), BOUNDARY + np.timedelta64(6, "h"), np.timedelta64(6, "h"))
    leads = ((valid - INIT) / np.timedelta64(1, "h")).astype(float)
    shape = (25, 1, valid.size, 1, 1)
    dims = ("number", "forecast_reference_time", "forecast_period", "latitude", "longitude")
    coords = {"number": np.arange(25), "forecast_reference_time": [INIT],
              "forecast_period": leads, "latitude": [10.0], "longitude": [105.0],
              "valid_time": ("forecast_period", valid)}
    return xr.Dataset({
        "t2m": (dims, np.full(shape, 300.0), {"units": "K"}),
        "sst": (dims, np.full(shape, 301.0), {"units": "K"}),
        "stl1": (dims, np.full(shape, 302.0), {"units": "K"}),
        "msl": (dims, np.full(shape, 101000.0), {"units": "Pa"}),
    }, coords=coords)


def accumulated_dataset():
    valid = np.arange(INIT + np.timedelta64(24, "h"), BOUNDARY + np.timedelta64(24, "h"), np.timedelta64(24, "h"))
    leads = ((valid - INIT) / np.timedelta64(1, "h")).astype(float)
    days = np.arange(1, valid.size + 1, dtype=float)
    values = np.broadcast_to(days[None, None, :, None, None], (25, 1, valid.size, 1, 1)).copy()
    dims = ("number", "forecast_reference_time", "forecast_period", "latitude", "longitude")
    coords = {"number": np.arange(25), "forecast_reference_time": [INIT],
              "forecast_period": leads, "latitude": [10.0], "longitude": [105.0],
              "valid_time": ("forecast_period", valid)}
    return xr.Dataset({
        "tp": (dims, values * 0.001, {"units": "m"}),
        "slhf": (dims, values * 86400 * -10, {"units": "J m**-2"}),
        "sshf": (dims, values * 86400 * -5, {"units": "J m**-2"}),
        "tsr": (dims, values * 86400 * 100, {"units": "J m**-2"}),
        "ttr": (dims, values * 86400 * -200, {"units": "J m**-2"}),
    }, coords=coords)


def test_monthly_build_keeps_members_boundaries_and_units():
    result = build_monthly_dataset(instant_dataset(), accumulated_dataset(), "1981-01")
    assert result.sizes == {"init_time": 1, "member": 25, "lead": 7, "latitude": 1, "longitude": 1}
    np.testing.assert_array_equal(result.lead, np.arange(7))
    np.testing.assert_allclose(result.t2m_mod, 26.85)
    np.testing.assert_allclose(result.msl, 1010.0)
    np.testing.assert_allclose(result.pr_mod.isel(lead=0), 31.0)
    np.testing.assert_allclose(result.pr_mod.isel(lead=1), 28.0)
    np.testing.assert_allclose(result.lh, -10.0)
    np.testing.assert_allclose(result.olr, 200.0)
    assert result.pr_mod.attrs["units"] == "mm month-1"
    assert result.t2m_mod.attrs["units"] == "degC"
    assert "no ensemble mean" in result.attrs["ensemble_processing"]


def test_monthly_build_promotes_valid_time_data_variable():
    instantaneous = instant_dataset().reset_coords("valid_time")
    accumulated = accumulated_dataset().reset_coords("valid_time")
    assert "valid_time" in instantaneous.data_vars
    result = build_monthly_dataset(instantaneous, accumulated, "1981-01")
    assert result.sizes["lead"] == 7


def test_missing_internal_six_hour_step_is_rejected():
    source = instant_dataset().t2m.isel(forecast_period=[0, *range(2, 848)])
    try:
        instantaneous_monthly(source, "1981-01", "t2m_mod", "degC", 1.0, -273.15)
    except ValueError as exc:
        assert "missing or duplicate" in str(exc)
    else:
        raise AssertionError("Missing internal 6-hour step was accepted")


def test_wrong_units_are_rejected():
    source = instant_dataset().t2m.copy()
    source.attrs["units"] = "degC"
    try:
        instantaneous_monthly(source, "1981-01", "t2m_mod", "degC", 1.0, -273.15)
    except ValueError as exc:
        assert "units must be K" in str(exc)
    else:
        raise AssertionError("Wrong source unit was accepted")



def test_run_validates_and_atomically_promotes(tmp_path):
    instantaneous = tmp_path / "instantaneous.nc"
    accumulated = tmp_path / "accumulated.nc"
    output = tmp_path / "monthly.nc"
    instant_dataset().to_netcdf(instantaneous)
    accumulated_dataset().to_netcdf(accumulated)
    target, report = run("1981-01", instantaneous, accumulated, output)
    assert target == output
    assert report["sizes"]["member"] == 25
    assert output.exists()
    assert not output.with_suffix(".nc.part").exists()