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

from dr_gcm.baseline import (
    _atomic,
    fit_monthly_climatology,
    fit_ridge,
    persistence_forecast,
    predict_monthly_climatology,
    predict_ridge,
    validate_model,
    validate_prediction,
)


def data():
    time = pd.date_range("2000-01-01", periods=72, freq="MS")
    index = np.arange(72, dtype=float)
    values = np.column_stack([index / 10, np.sin(index / 3), np.ones(72)])
    predictors = xr.DataArray(
        values, dims=("time", "feature"),
        coords={"time": time, "feature": ["trend", "oscillation", "constant"]},
        name="predictors",
    )
    target_values = (5 + 2 * values[:, 0] - 3 * values[:, 1])[:, None]
    targets = xr.DataArray(
        target_values, dims=("time", "target"),
        coords={"time": time, "target": ["temperature"]}, name="targets",
    )
    return predictors, targets


def test_linear_recovers_relation_and_constant_feature_is_zero():
    predictors, targets = data()
    model = fit_ridge(predictors, targets, 0.0, "2000-01-01", "2003-12-31")
    prediction = predict_ridge(predictors, model)
    np.testing.assert_allclose(prediction, targets, atol=1e-10)
    assert model.attrs["model"] == "LinearRegression"
    assert model.attrs["training_samples"] == 48
    assert int(model.constant_feature.sel(feature="constant")) == 1
    assert float(model.coefficient.sel(feature="constant")) == 0


def test_scaler_uses_training_slice_only_and_ridge_shrinks():
    predictors, targets = data()
    linear = fit_ridge(predictors, targets, 0.0, "2000-01-01", "2003-12-31")
    modified = predictors.copy()
    modified.loc[{"time": slice("2004-01-01", None), "feature": "trend"}] = 1_000_000
    refit = fit_ridge(modified, targets, 10.0, "2000-01-01", "2003-12-31")
    ridge = fit_ridge(predictors, targets, 10.0, "2000-01-01", "2003-12-31")
    xr.testing.assert_equal(refit.feature_mean, ridge.feature_mean)
    assert np.linalg.norm(ridge.coefficient) < np.linalg.norm(linear.coefficient)
    assert ridge.attrs["scaler_fit_scope"] == "training slice only"


def test_missing_training_value_is_rejected():
    predictors, targets = data()
    predictors[0, 0] = np.nan
    try:
        fit_ridge(predictors, targets, 1.0, "2000-01-01", "2003-12-31")
    except ValueError as exc:
        assert "non-finite" in str(exc)
    else:
        raise AssertionError("Missing training predictor was accepted")


def test_model_and_prediction_atomic_outputs(tmp_path):
    predictors, targets = data()
    model = fit_ridge(predictors, targets, 1.0, "2000-01-01", "2003-12-31")
    model_path = tmp_path / "ridge.nc"
    _, model_report = _atomic(model, model_path, validate_model)
    assert model_report["sizes"]["feature"] == 3
    prediction = predict_ridge(predictors, model).rename("t2m_prediction")
    output = tmp_path / "prediction.nc"
    _, report = _atomic(
        prediction.to_dataset(), output,
        lambda path: validate_prediction(path, "t2m_prediction"),
    )
    assert report["sizes"]["time"] == 72
    assert not output.with_suffix(".nc.part").exists()


def test_monthly_climatology_is_fit_only_on_training_slice():
    time = pd.date_range("2000-01-01", periods=36, freq="MS")
    month = np.tile(np.arange(1, 13, dtype=float), 3)
    year_offset = np.repeat([0.0, 2.0, 1000.0], 12)
    observed = xr.DataArray(month + year_offset, dims="time", coords={"time": time}, attrs={"units": "degC"})
    model = fit_monthly_climatology(observed, "2000-01-01", "2001-12-31")
    np.testing.assert_allclose(model.climatology, np.arange(1, 13) + 1.0)
    assert model.attrs["training_samples"] == 24
    valid = xr.DataArray(
        pd.to_datetime(["2003-01-01", "2003-02-01"]),
        dims="lead", coords={"lead": [1, 2]},
    )
    prediction = predict_monthly_climatology(valid, model)
    np.testing.assert_allclose(prediction, [2.0, 3.0])
    np.testing.assert_array_equal(prediction.valid_time, valid)


def test_persistence_uses_exact_month_before_initialization():
    observed = xr.DataArray(
        np.arange(12, dtype=float), dims="time",
        coords={"time": pd.date_range("2024-05-01", periods=12, freq="MS")},
        attrs={"units": "mm month-1"},
    )
    valid = xr.DataArray(
        pd.date_range("2025-05-01", periods=7, freq="MS"),
        dims="lead", coords={"lead": np.arange(7)},
    )
    result = persistence_forecast(observed, "2025-05", valid)
    np.testing.assert_allclose(result, np.full(7, 11.0))
    assert result.attrs["source_valid_time"].startswith("2025-04")
    try:
        persistence_forecast(observed.isel(time=slice(None, -1)), "2025-05", valid)
    except ValueError as exc:
        assert "exactly one" in str(exc)
    else:
        raise AssertionError("Persistence accepted missing previous-month observation")