import json

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

from dr_gcm.temporal_cv import build_plan, rolling_origin_folds, write_plan


def test_expanding_rolling_origin_folds_are_strictly_ordered():
    sample_years = np.repeat(np.arange(1981, 2021), 12)
    folds = rolling_origin_folds(sample_years, 20, 5, 5, 5)
    assert len(folds) == 3
    assert folds[0].train_years[[0, -1]].tolist() == [1981, 2000]
    assert folds[0].validation_years[[0, -1]].tolist() == [2001, 2005]
    assert folds[0].test_years[[0, -1]].tolist() == [2006, 2010]
    assert folds[-1].test_years[[0, -1]].tolist() == [2016, 2020]
    for fold in folds:
        masks = fold.masks(sample_years)
        assert not np.any(masks["train"] & masks["validation"])
        assert not np.any(masks["train"] & masks["test"])
        assert not np.any(masks["validation"] & masks["test"])
        assert fold.train_years.max() < fold.validation_years.min()
        assert fold.validation_years.max() < fold.test_years.min()


def test_gap_years_are_excluded_between_blocks():
    years = np.arange(1981, 2021)
    fold = rolling_origin_folds(years, 20, 4, 4, 4, gap_years=1)[0]
    assert fold.train_years[-1] == 2000
    assert fold.validation_years[0] == 2002
    assert fold.validation_years[-1] == 2005
    assert fold.test_years[0] == 2007


def test_missing_available_year_is_rejected():
    years = np.array([1981, 1982, 1984, 1985, 1986, 1987])
    try:
        rolling_origin_folds(years, 2, 1, 1, 1)
    except ValueError as exc:
        assert "contiguous" in str(exc)
    else:
        raise AssertionError("Missing year was accepted")


def test_plan_uses_real_time_coordinate_and_atomic_json(tmp_path):
    source = tmp_path / "monthly.nc"
    output = tmp_path / "cv_plan.json"
    time = pd.date_range("1981-01-01", "2020-12-01", freq="MS")
    xr.Dataset({"value": ("time", np.zeros(time.size))}, coords={"time": time}).to_netcdf(source)
    plan = build_plan(source, "time", 20, 5, 5, 5)
    assert plan["available_year_range"] == [1981, 2020]
    assert len(plan["folds"]) == 3
    assert plan["folds"][0]["sample_counts"] == {
        "train": 240, "validation": 60, "test": 60
    }
    write_plan(plan, output)
    loaded = json.loads(output.read_text(encoding="utf-8"))
    assert loaded["settings"]["min_train_years"] == 20
    assert not output.with_suffix(".json.part").exists()