import numpy as np
import pandas as pd
import xarray as xr

from dr_gcm.operational import combine_history_forecast, run


def inputs():
    observed = xr.DataArray(
        np.arange(12, dtype=float)[:, None, None],
        dims=("time", "latitude", "longitude"),
        coords={"time": pd.date_range("2024-05-01", periods=12, freq="MS"),
                "latitude": [16.0], "longitude": [106.0]},
        name="pr_obs", attrs={"units": "mm month-1"},
    )
    valid = pd.date_range("2025-05-01", periods=7, freq="MS")
    forecast = xr.DataArray(
        np.arange(14, dtype=float).reshape(2, 7, 1, 1),
        dims=("member", "lead", "latitude", "longitude"),
        coords={"member": [0, 1], "lead": np.arange(7),
                "valid_time": ("lead", valid),
                "latitude": [16.0], "longitude": [106.0]},
        name="pr_mod", attrs={"units": "mm month-1"},
    )
    return observed, forecast


def test_history_is_broadcast_and_lt0_to_lt6_are_retained():
    observed, forecast = inputs()
    result = combine_history_forecast(observed, forecast, "2025-05")
    assert result.sizes == {"member": 2, "time": 12, "latitude": 1, "longitude": 1}
    np.testing.assert_array_equal(
        result.time.values.astype("datetime64[M]"),
        np.arange(np.datetime64("2024-12"), np.datetime64("2025-12"), dtype="datetime64[M]"),
    )
    np.testing.assert_allclose(result.isel(member=0, time=slice(0, 5)).values[:, 0, 0], [7, 8, 9, 10, 11])
    np.testing.assert_allclose(result.isel(member=1, time=slice(0, 5)).values[:, 0, 0], [7, 8, 9, 10, 11])
    np.testing.assert_allclose(result.isel(member=0, time=slice(5, None)).values[:, 0, 0], np.arange(7))
    assert result.attrs["history_months"] == 5


def test_missing_required_observed_month_is_rejected():
    observed, forecast = inputs()
    observed = observed.isel(time=slice(None, -1))
    try:
        combine_history_forecast(observed, forecast, "2025-05")
    except ValueError as exc:
        assert "exactly one value" in str(exc)
    else:
        raise AssertionError("Missing observed history was accepted")


def test_wrong_forecast_valid_time_is_rejected():
    observed, forecast = inputs()
    forecast = forecast.assign_coords(valid_time=("lead", pd.date_range("2025-06-01", periods=7, freq="MS")))
    try:
        combine_history_forecast(observed, forecast, "2025-05")
    except ValueError as exc:
        assert "does not match" in str(exc)
    else:
        raise AssertionError("Shifted forecast valid_time was accepted")


def test_cli_run_writes_validated_atomic_output(tmp_path):
    observed, forecast = inputs()
    observed_path = tmp_path / "observed.nc"
    forecast_path = tmp_path / "forecast.nc"
    output = tmp_path / "combined.nc"
    observed.to_dataset().to_netcdf(observed_path)
    forecast.to_dataset().to_netcdf(forecast_path)
    target, report = run(
        observed_path, forecast_path, output, "2025-05",
        "pr_obs", "pr_mod", "pr", history_months=5,
    )
    assert target == output
    assert report["sizes"]["member"] == 2
    assert not output.with_suffix(".nc.part").exists()