import numpy as np
import pytest
import xarray as xr

from dr_gcm.climatology import climatology_periods, model_climatology


def _write_model(path, initializations):
    init = np.asarray(initializations, dtype="datetime64[ns]")
    member = np.array([0, 1])
    lead = np.arange(7)
    values = np.empty((init.size, member.size, lead.size, 1, 1), dtype=float)
    for index, timestamp in enumerate(init):
        year_offset = (int(str(timestamp)[:4]) - 2000) * 10
        values[index, :, :, 0, 0] = (
            year_offset + member[:, None] * 100 + lead[None, :]
        )
    xr.Dataset(
        {
            "t2m_mod": (
                ("init_time", "member", "lead", "latitude", "longitude"),
                values,
                {"units": "degC"},
            )
        },
        coords={
            "init_time": init,
            "member": member,
            "lead": lead,
            "latitude": [20.0],
            "longitude": [105.0],
        },
    ).to_netcdf(path)


def test_model_climatology_retains_members_and_excludes_lead_zero(tmp_path):
    months = [
        f"{year}-{month:02d}-01" for year in (2000, 2001) for month in range(1, 13)
    ]
    _write_model(tmp_path / "monthly.nc", months)
    result = model_climatology(
        str(tmp_path / "*.nc"), "t2m_mod", "2000-01-01", "2001-12-31"
    )
    assert result.sizes == {
        "longitude": 1, "latitude": 1, "member": 2, "lead": 6, "month": 12
    }
    assert result.lead.values.tolist() == [1, 2, 3, 4, 5, 6]
    np.testing.assert_allclose(result.sel(member=0, lead=3), 8.0)
    np.testing.assert_allclose(result.sel(member=1, lead=3), 108.0)
    assert result.attrs["ensemble_processing"] == "members retained; no ensemble mean"
    assert result.attrs["initialization_month_count"] == 24


def test_model_climatology_rejects_duplicate_initialization_months(tmp_path):
    _write_model(tmp_path / "first.nc", ["2000-01-01"])
    _write_model(tmp_path / "second.nc", ["2000-01-01"])
    with pytest.raises(ValueError, match="Duplicate model initialization month"):
        model_climatology(str(tmp_path / "*.nc"), "t2m_mod", "2000", "2000-12-31")


def test_climatology_periods_separate_observations_from_model_initializations():
    config = {
        "baseline": {
            "observation": {"start": "1981-01-01", "end": "2021-06-30"},
            "model_initialization": {"start": "1981-01-01", "end": "2020-12-31"},
        }
    }
    assert climatology_periods(config) == (
        "1981-01-01",
        "2021-06-30",
        "1981-01-01",
        "2020-12-31",
    )


def test_climatology_periods_support_legacy_flat_config():
    config = {"baseline": {"start": "1981-01-01", "end": "2020-12-31"}}
    assert climatology_periods(config) == (
        "1981-01-01",
        "2020-12-31",
        "1981-01-01",
        "2020-12-31",
    )