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

from dr_gcm.spei import accumulated_balance, fit_loglogistic, fit_thornthwaite_parameters, spei_from_parameters, thornthwaite_pet, water_balance


def monthly_field(values):
    return xr.DataArray(
        values[:, None, None],
        dims=("time", "latitude", "longitude"),
        coords={"time": pd.date_range("1981-01-01", periods=len(values), freq="MS"), "latitude": [16.0], "longitude": [106.0]},
    )


def test_thornthwaite_pet_is_nonnegative_and_zero_for_freezing_months():
    temperature = monthly_field(np.array([-1.0, 0.0, 22.0, 26.0] * 3))
    pet = thornthwaite_pet(temperature, fit_thornthwaite_parameters(temperature))
    assert pet.name == "pet"
    assert float(pet.isel(time=0)) == 0
    assert float(pet.isel(time=2)) > 0


def test_accumulation_and_saved_parameters_produce_finite_spei():
    values = np.arange(60, dtype=float) * 3 - 90
    balance = monthly_field(values)
    assert np.isnan(accumulated_balance(balance, 3).isel(time=1)).all()
    assert float(accumulated_balance(balance, 3).isel(time=2)) == -261
    params = fit_loglogistic(balance, scales=(1, 3))
    spei = spei_from_parameters(balance, params, 3)
    assert spei.name == "spei_3"
    assert np.isfinite(spei.isel(time=slice(26, None))).all()


def test_water_balance_uses_precipitation_minus_pet():
    temperature = monthly_field(np.full(12, 24.0))
    result = water_balance(monthly_field(np.full(12, 200.0)), temperature, fit_thornthwaite_parameters(temperature))
    assert result.name == "water_balance"
    assert float(result.mean()) < 200


def test_water_balance_clips_negative_precipitation_before_pet():
    temperature = monthly_field(np.zeros(12))
    precipitation = monthly_field(np.array([-5.0, 10.0] * 6))
    result = water_balance(precipitation, temperature, fit_thornthwaite_parameters(temperature))
    np.testing.assert_allclose(result.values[:, 0, 0], [0.0, 10.0] * 6)
    assert "below zero" in result.attrs["precipitation_negative_value_policy"]



def test_thornthwaite_parameters_are_climatological_and_reusable():
    calibration = monthly_field(np.tile(np.arange(15.0, 27.0), 2))
    parameters = fit_thornthwaite_parameters(calibration)
    assert "time" not in parameters.dims
    assert "year" not in parameters.dims
    assert parameters.attrs["operational_rule"].startswith("reuse saved")
    future = xr.DataArray(
        np.full((12, 1, 1), 25.0),
        dims=("time", "latitude", "longitude"),
        coords={"time": __import__("pandas").date_range("2030-01-01", periods=12, freq="MS"),
                "latitude": [16.0], "longitude": [106.0]},
    )
    pet = thornthwaite_pet(future, parameters)
    assert np.isfinite(pet).all()
    assert pet.attrs["parameter_calibration_start"].startswith("1981")