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

from dr_gcm.spei_workflow import (
    apply_parameters,
    fit_parameters,
    validate_output,
    validate_parameters,
    write_atomic,
)


def fields(periods=60):
    time = pd.date_range("1981-01-01", periods=periods, freq="MS")
    index = np.arange(periods, dtype=float)
    seasonal = np.sin(2 * np.pi * index / 12)
    pr = xr.DataArray(
        (120 + 30 * seasonal + index * 0.7)[:, None, None],
        dims=("time", "latitude", "longitude"),
        coords={"time": time, "latitude": [16.0], "longitude": [106.0]},
        attrs={"units": "mm month-1"},
    )
    t2m = xr.DataArray(
        (24 + 4 * seasonal + index * 0.02)[:, None, None],
        dims=pr.dims,
        coords=pr.coords,
        attrs={"units": "degC"},
    )
    return pr, t2m


def test_fit_and_apply_reusable_parameters():
    pr, t2m = fields()
    parameters = fit_parameters(pr, t2m, "1981-01-01", "1985-12-31")
    assert set(parameters.data_vars) == {
        "heat_index", "thornthwaite_exponent", "loglogistic_parameters"
    }
    np.testing.assert_array_equal(parameters.timescale, [1, 3, 6])
    np.testing.assert_array_equal(parameters.month, np.arange(1, 13))
    result = apply_parameters(pr, t2m, parameters)
    assert set(result.data_vars) == {"spei_1", "spei_3", "spei_6"}
    assert np.isfinite(result.spei_1.isel(time=slice(12, None))).all()
    assert result.attrs["parameter_calibration_start"].startswith("1981")


def test_monthly_gap_and_wrong_units_are_rejected():
    pr, t2m = fields()
    try:
        fit_parameters(pr.isel(time=[0, *range(2, 60)]), t2m.isel(time=[0, *range(2, 60)]))
    except ValueError as exc:
        assert "missing monthly" in str(exc)
    else:
        raise AssertionError("Monthly gap was accepted")
    pr.attrs["units"] = "m"
    try:
        fit_parameters(pr, t2m)
    except ValueError as exc:
        assert "monthly mm" in str(exc)
    else:
        raise AssertionError("Wrong precipitation units were accepted")


def test_parameter_and_spei_outputs_are_atomic_and_validated(tmp_path):
    pr, t2m = fields()
    parameters = fit_parameters(pr, t2m, "1981-01-01", "1985-12-31")
    parameter_path = tmp_path / "parameters.nc"
    _, parameter_report = write_atomic(parameters, parameter_path, validate_parameters)
    assert parameter_report["sizes"]["timescale"] == 3
    assert not parameter_path.with_suffix(".nc.part").exists()
    spei = apply_parameters(pr, t2m, parameters)
    output = tmp_path / "spei.nc"
    _, output_report = write_atomic(spei, output, validate_output)
    assert output_report["variables"] == ["spei_1", "spei_3", "spei_6"]
    assert not output.with_suffix(".nc.part").exists()