import numpy as np
import xarray as xr

from dr_gcm.drought import classify_spei, drought_probabilities, run


def test_classification_respects_all_confirmed_boundaries():
    values = xr.DataArray(
        [-2.1, -2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0, np.nan],
        dims="point",
    )
    result = classify_spei(values)
    np.testing.assert_allclose(
        result.values,
        [-4, -3, -3, -2, -1, 0, 1, 2, 3, 4, np.nan],
        equal_nan=True,
    )
    assert result.attrs["flag_meanings"] == "D4 D3 D2 D1 D0 W1 W2 W3 W4"


def test_nested_probabilities_use_all_supplied_components_and_strict_d4():
    components = xr.DataArray(
        [-0.25, -0.5, -1.0, -1.5, -2.0, -2.1],
        dims="component",
    )
    result = drought_probabilities(components)
    np.testing.assert_allclose(
        [result.probability_d1, result.probability_d2,
         result.probability_d3, result.probability_d4],
        [5 / 6, 4 / 6, 3 / 6, 1 / 6],
    )
    assert int(result.valid_component_count) == 6


def test_all_missing_is_missing_but_partial_missing_is_rejected():
    all_missing = xr.DataArray([np.nan, np.nan], dims="component")
    result = drought_probabilities(all_missing)
    assert np.isnan(result.probability_d1)
    partial = xr.DataArray([np.nan, -1.0], dims="component")
    try:
        drought_probabilities(partial)
    except ValueError as exc:
        assert "Partial missing" in str(exc)
    else:
        raise AssertionError("Partial missing ensemble was accepted")


def test_cli_run_writes_validated_atomic_product(tmp_path):
    source = tmp_path / "spei.nc"
    output = tmp_path / "drought.nc"
    xr.Dataset(
        {
            "spei_ens": ("lead", [-0.6, -2.1]),
            "spei_components": (
                ("component", "lead"),
                [[-0.6, -2.1], [-0.4, -1.6], [-1.1, -2.2]],
            ),
        },
        coords={"component": ["member_a", "member_b", "spei_ens"], "lead": [1, 2]},
    ).to_netcdf(source)
    target, report = run(source, output)
    assert target == output
    assert output.exists()
    assert not output.with_suffix(".nc.part").exists()
    assert "probability_d4" in report["variables"]
    with xr.open_dataset(output) as result:
        np.testing.assert_allclose(result.probability_d1, [2 / 3, 1.0])
        np.testing.assert_allclose(result.probability_d4, [0.0, 2 / 3])