# SPDX-FileCopyrightText: 2025 Contributors to the OpenSTEF project <short.term.energy.forecasts@alliander.com>
#
# SPDX-License-Identifier: MPL-2.0

from datetime import timedelta

import pytest

from openstef_core.datasets.validated_datasets import EnsembleForecastDataset
from openstef_core.exceptions import NotFittedError
from openstef_core.types import LeadTime, Q
from openstef_meta.models.forecast_combiners.learned_weights_combiner import (
    LGBMCombinerHyperParams,
    LogisticCombinerHyperParams,
    RFCombinerHyperParams,
    WeightsCombiner,
    XGBCombinerHyperParams,
)


@pytest.fixture(params=["lgbm", "xgboost", "rf", "logistic"])
def classifier(request: pytest.FixtureRequest) -> str:
    """Fixture to provide different classifier types for LearnedWeightsCombiner tests."""
    return request.param


@pytest.fixture
def combiner(classifier: str) -> WeightsCombiner:
    """Fixture to create a WeightsCombiner based on the classifier type."""
    if classifier == "lgbm":
        hp = LGBMCombinerHyperParams(n_leaves=5, n_estimators=10)
    elif classifier == "xgboost":
        hp = XGBCombinerHyperParams(n_estimators=10)
    elif classifier == "rf":
        hp = RFCombinerHyperParams(n_estimators=10, n_leaves=5)
    elif classifier == "logistic":
        hp = LogisticCombinerHyperParams()
    else:
        msg = f"Unsupported classifier type: {classifier}"
        raise ValueError(msg)

    return WeightsCombiner(hyperparams=hp, quantiles=[Q(0.1), Q(0.5), Q(0.9)], horizons=[LeadTime(timedelta(days=1))])


def test_quantile_weights_combiner__fit_predict(
    ensemble_dataset: EnsembleForecastDataset,
    combiner: WeightsCombiner,
):
    """Test basic fit and predict workflow with comprehensive output validation."""
    # Arrange
    expected_quantiles = combiner.quantiles

    # Act
    combiner.fit(ensemble_dataset)
    result = combiner.predict(ensemble_dataset)

    # Assert
    assert combiner.is_fitted, "Model should be fitted after calling fit()"

    expected_columns = [q.format() for q in expected_quantiles]
    expected_columns.append("load")
    assert list(result.data.columns) == expected_columns, (
        f"Expected columns {expected_columns}, got {list(result.data.columns)}"
    )

    assert not result.data.isna().any().any(), "Forecast should not contain NaN or None values"

    stds = result.data.std()
    assert (stds > 0).all(), f"All columns should have variation, got stds: {dict(stds)}"


def test_weights_combiner_not_fitted_error(
    ensemble_dataset: EnsembleForecastDataset,
    combiner: WeightsCombiner,
):
    """Test that NotFittedError is raised when predicting before fitting."""
    with pytest.raises(NotFittedError):
        combiner.predict(ensemble_dataset)
