import numpy as np
import pytest

from skchange.costs import VALID_SAVINGS, saving_factory
from skchange.datasets.generate import generate_alternating_data


@pytest.mark.parametrize("savings", VALID_SAVINGS)
def test_savings(savings):
    """Test all available savings."""
    n = 50
    df = generate_alternating_data(n_segments=1, segment_length=n, p=1, random_state=5)
    savings_f, init_savings_f = saving_factory(savings)
    params = init_savings_f(df.values)
    savings_values = np.zeros(n)
    starts = np.arange(n - 10)
    ends = np.repeat(n - 1, len(starts))
    savings_values = savings_f(params, starts, ends=ends)
    assert np.all(savings_values >= 0.0)


def test_custom_savings():
    """Test custom savings."""

    def init_savings_f(X: np.ndarray) -> np.ndarray:
        return X

    def savings_f(params: np.ndarray, start: int, end: int, split: int) -> float:
        return 10.0

    assert init_savings_f(np.zeros(1)) == np.zeros(1)
    assert savings_f(np.zeros(1), 0, 1, 0) == 10.0
    assert (savings_f, init_savings_f) == saving_factory((savings_f, init_savings_f))


def test_savings_factory_raises_on_invalid_savings():
    """Test that savings_factory raises an error on invalid savings."""
    with pytest.raises(ValueError):
        saving_factory("invalid_savings")
