import numpy as np
import xarray as xr

from dr_gcm.verification import continuous_metrics, contingency_metrics, spei_drought_event


def test_temperature_metrics_and_paired_missing_samples():
    forecast = xr.DataArray([2.0, 4.0, 6.0, np.nan], dims="sample")
    observed = xr.DataArray([1.0, 4.0, 8.0, 99.0], dims="sample")
    result = continuous_metrics(forecast, observed, "sample")
    np.testing.assert_allclose(result.me, -1 / 3)
    np.testing.assert_allclose(result.mae, 1.0)
    np.testing.assert_allclose(result.rmse, np.sqrt(5 / 3))
    np.testing.assert_allclose(result.correlation, np.corrcoef([2, 4, 6], [1, 4, 8])[0, 1])
    assert int(result.sample_count) == 3
    assert "rmae" not in result


def test_precipitation_rmae_matches_confirmed_definition():
    forecast = xr.DataArray([2.0, 6.0], dims="sample")
    observed = xr.DataArray([1.0, 3.0], dims="sample")
    result = continuous_metrics(forecast, observed, "sample", precipitation=True)
    np.testing.assert_allclose(result.me, 2.0)
    np.testing.assert_allclose(result.rmae, 100.0)
    assert "mae" not in result
    assert "rmse" not in result


def test_contingency_table_and_rare_event_metrics():
    forecast = xr.DataArray([1, 1, 0, 0], dims="sample")
    observed = xr.DataArray([1, 0, 1, 0], dims="sample")
    result = contingency_metrics(forecast, observed, "sample")
    for name in ("hits", "misses", "false_alarms", "correct_negatives"):
        assert int(result[name]) == 1
    np.testing.assert_allclose(result.pod, 0.5)
    np.testing.assert_allclose(result.pofd, 0.5)
    np.testing.assert_allclose(result.far, 0.5)
    np.testing.assert_allclose(result.fbi, 1.0)
    np.testing.assert_allclose(result.csi, 1 / 3)
    np.testing.assert_allclose(result.edi, 0.0)
    np.testing.assert_allclose(result.sedi, 0.0)


def test_drought_events_use_strict_d4_boundary():
    spei = xr.DataArray([-2.1, -2.0, -1.5, -0.5, np.nan], dims="sample")
    np.testing.assert_allclose(
        spei_drought_event(spei, 4), [1, 0, 0, 0, np.nan], equal_nan=True
    )
    np.testing.assert_allclose(
        spei_drought_event(spei, 3), [1, 1, 1, 0, np.nan], equal_nan=True
    )


def test_undefined_edi_sedi_are_missing_without_edge_correction():
    result = contingency_metrics(
        xr.DataArray([1, 0], dims="sample"),
        xr.DataArray([1, 0], dims="sample"),
        "sample",
    )
    assert np.isnan(result.edi)
    assert np.isnan(result.sedi)