"""Tests for elastic net / LASSO / Ridge regression module."""

import pytest
from horizon._horizon import (
    ElasticNetResult,
    elastic_net_fit,
    elastic_net_predict,
    lasso_fit,
    ridge_fit,
    elastic_net_cv,
)


def _make_linear_data(n=100, p=5, noise=0.1, seed=42):
    """Generate linear data: y = x[0] + 2*x[1] + noise."""
    # Simple deterministic pseudo-random
    x = []
    y = []
    state = seed
    for i in range(n):
        row = []
        for j in range(p):
            state = (state * 1103515245 + 12345) & 0x7FFFFFFF
            row.append((state / 0x7FFFFFFF) * 2.0 - 1.0)
        x.append(row)
        state = (state * 1103515245 + 12345) & 0x7FFFFFFF
        eps = ((state / 0x7FFFFFFF) * 2.0 - 1.0) * noise
        y.append(1.0 * row[0] + 2.0 * row[1] + eps)
    return x, y


class TestElasticNetFit:
    def test_basic_fit(self):
        x, y = _make_linear_data()
        result = elastic_net_fit(x, y, alpha=0.01, l1_ratio=0.5)
        assert isinstance(result, ElasticNetResult)
        assert len(result.coefficients) == 5
        assert result.r_squared > 0.5

    def test_lasso_sparsity(self):
        x, y = _make_linear_data(p=10)
        result = elastic_net_fit(x, y, alpha=0.1, l1_ratio=1.0)
        # LASSO should zero out some coefficients
        assert result.nonzero_count < 10

    def test_ridge_no_sparsity(self):
        x, y = _make_linear_data(p=5)
        result = elastic_net_fit(x, y, alpha=0.01, l1_ratio=0.0)
        # Ridge keeps all coefficients nonzero (usually)
        assert result.nonzero_count >= 2

    def test_repr(self):
        x, y = _make_linear_data(n=50)
        result = elastic_net_fit(x, y)
        assert "ElasticNetResult" in repr(result)

    def test_negative_alpha(self):
        x, y = _make_linear_data(n=20)
        with pytest.raises(ValueError):
            elastic_net_fit(x, y, alpha=-1.0)

    def test_empty_input(self):
        with pytest.raises(ValueError):
            elastic_net_fit([], [])

    def test_mismatched_lengths(self):
        with pytest.raises(ValueError):
            elastic_net_fit([[1.0, 2.0]], [1.0, 2.0])

    def test_single_row(self):
        # Should work with 1 sample (though R² may be 0)
        result = elastic_net_fit([[1.0, 2.0]], [3.0])
        assert len(result.coefficients) == 2


class TestElasticNetPredict:
    def test_predict(self):
        x, y = _make_linear_data(n=50)
        result = elastic_net_fit(x, y, alpha=0.01)
        preds = elastic_net_predict(x, result)
        assert len(preds) == 50
        # Predictions should be close to actual
        errors = [abs(p - a) for p, a in zip(preds, y)]
        assert sum(errors) / len(errors) < 2.0

    def test_dimension_mismatch(self):
        result = elastic_net_fit([[1.0, 2.0]], [3.0])
        with pytest.raises(ValueError):
            elastic_net_predict([[1.0]], result)  # Wrong feature count


class TestLassoFit:
    def test_lasso_shortcut(self):
        x, y = _make_linear_data()
        result = lasso_fit(x, y, alpha=0.1)
        assert isinstance(result, ElasticNetResult)
        assert result.l1_ratio == 1.0


class TestRidgeFit:
    def test_ridge_shortcut(self):
        x, y = _make_linear_data()
        result = ridge_fit(x, y, alpha=0.1)
        assert isinstance(result, ElasticNetResult)
        assert result.l1_ratio == 0.0


class TestElasticNetCV:
    def test_cv_basic(self):
        x, y = _make_linear_data(n=50, p=3)
        result = elastic_net_cv(x, y, n_folds=3)
        assert isinstance(result, ElasticNetResult)
        assert result.alpha > 0.0
        assert result.r_squared >= -1.0  # Can be negative for bad fits

    def test_cv_custom_alphas(self):
        x, y = _make_linear_data(n=50)
        result = elastic_net_cv(x, y, alphas=[0.01, 0.1], l1_ratios=[0.5, 1.0])
        assert result.alpha in [0.01, 0.1]
        assert result.l1_ratio in [0.5, 1.0]

    def test_cv_small_data(self):
        x, y = _make_linear_data(n=10, p=2)
        result = elastic_net_cv(x, y, n_folds=2)
        assert isinstance(result, ElasticNetResult)
