import numpy as np
import pytest

from sgd_classifier import SUPPORTED_LOSSES, SGDClassifier


def dataset(seed=4, samples=61, features=8):
    rng = np.random.default_rng(seed)
    X = rng.normal(size=(samples, features))
    y = np.where(X[:, 0] - 0.5 * X[:, 1] > 0, "positive", "negative")
    return X, y


@pytest.mark.parametrize("loss", SUPPORTED_LOSSES)
@pytest.mark.parametrize("selection", ["random", "permutation"])
def test_every_loss_and_selection_mode_trains(loss, selection):
    X, y = dataset()
    model = SGDClassifier(
        loss=loss, batch_size=13, batch_selection=selection, max_epochs=2,
        random_state=7,
    ).fit(X, y)
    assert model.predict(X).shape == (61,)
    assert model.n_updates_ == 10


def test_train_step_updates_and_returns_model():
    X, y = dataset(samples=20)
    model = SGDClassifier(loss="log_loss", learning_rate=0.1)
    returned = model.train_step(X, y)
    assert returned is model
    assert model.n_updates_ == 1
    assert np.any(model.coef_ != 0)


def test_get_update_previews_without_applying_it():
    X, y = dataset(samples=20)
    model = SGDClassifier(loss="log_loss", learning_rate=0.1)
    coef_update, intercept_update = model.get_update(X, y)
    np.testing.assert_array_equal(model.coef_, np.zeros_like(model.coef_))
    np.testing.assert_array_equal(model.intercept_, np.zeros_like(model.intercept_))
    assert model.n_updates_ == 0
    model.train_step(X, y)
    np.testing.assert_allclose(model.coef_, coef_update)
    np.testing.assert_allclose(model.intercept_, intercept_update)


def test_evaluate_returns_dataset_loss_and_accuracy():
    X, y = dataset(samples=30)
    model = SGDClassifier(loss="log_loss").fit(X, y)
    metrics = model.evaluate(X, y)
    assert metrics["loss"] >= 0
    assert 0 <= metrics["accuracy"] <= 1


def test_train_step_can_declare_classes_missing_from_batch():
    model = SGDClassifier()
    model.train_step([[1.0, 2.0]], ["a"], classes=["a", "b"])
    assert model.classes_.tolist() == ["a", "b"]


def test_permutation_batches_cover_each_example_once():
    model = SGDClassifier(batch_size=4, batch_selection="permutation")
    batches = list(model._batches(11, np.random.default_rng(1)))
    assert sorted(np.concatenate(batches).tolist()) == list(range(11))
    assert [len(batch) for batch in batches] == [4, 4, 3]


def test_seed_is_reproducible():
    X, y = dataset(samples=50)
    options = dict(batch_selection="random", max_epochs=3, random_state=9)
    first = SGDClassifier(**options).fit(X, y)
    second = SGDClassifier(**options).fit(X, y)
    np.testing.assert_array_equal(first.coef_, second.coef_)