"""timecast.learners.AR: testing"""
import flax
import jax
import jax.numpy as jnp
import numpy as onp
import pytest
from sklearn.linear_model import Ridge

from timecast.learners import AR
from timecast.learners._ar import _compute_kernel_bias
from timecast.utils import historify
from timecast.utils import random


shapes = [(32, 4), (4, 32), (10,), (10, 1), (1,), (1, 10), (10, 10)]


def create_ar(shape, output_dim, history_len, history=None):
    """Create AR model"""
    model_def = AR.partial(output_dim=output_dim, history_len=history_len, history=history)
    with flax.nn.stateful() as state:
        _, params = model_def.init_by_shape(jax.random.PRNGKey(0), [shape])
        model = flax.nn.Model(model_def, params)
    return model, state


@pytest.mark.parametrize("shape", shapes)
@pytest.mark.parametrize("output_dim", [1, 10])
@pytest.mark.parametrize("history_len", [1, 10])
def test_ar_ys_shape(shape, output_dim, history_len):
    """Test output shape"""
    model, state = create_ar(shape, output_dim, history_len)
    with flax.nn.stateful(state) as _:
        ys = model(jax.random.uniform(random.generate_key(), shape=shape))

    assert ys.shape == (output_dim,)


@pytest.mark.parametrize("shape", shapes)
@pytest.mark.parametrize("output_dim", [1, 10])
@pytest.mark.parametrize("history_len", [1, 10])
def test_ar_history_shape(shape, output_dim, history_len):
    """Test history shape"""
    model, state = create_ar(shape, output_dim, history_len)
    with flax.nn.stateful(state) as state:
        _ = model(jax.random.uniform(random.generate_key(), shape=shape))

    print(state.as_dict())

    input_features = shape[0] if len(shape) == 1 else shape[1]
    assert state.as_dict()["/"]["history"].shape == (history_len, input_features)


@pytest.mark.parametrize("shape", shapes)
@pytest.mark.parametrize("output_dim", [1, 10])
@pytest.mark.parametrize("history_len", [1, 10])
def test_ar_history(shape, output_dim, history_len):
    """Test history values"""
    input_features = shape[0] if len(shape) == 1 else shape[1]
    history = jax.random.uniform(random.generate_key(), shape=(history_len, input_features))
    model, state = create_ar(shape, output_dim, history_len, history)

    onp.testing.assert_array_almost_equal(history, state.as_dict()["/"]["history"])


@pytest.mark.parametrize("shape", shapes)
@pytest.mark.parametrize("fit_intercept", [True, False])
@pytest.mark.parametrize("alpha", [0.0, 0.5, 1.0])
def test_compute_kernel_bias(shape, fit_intercept, alpha):
    """Test kernel and bias computation"""
    if len(shape) == 1:
        shape += (1,)
    X = jax.random.uniform(random.generate_key(), shape=shape)
    Y = jax.random.uniform(random.generate_key(), shape=(shape[0],))

    # Ignore underdetermined systems for now
    if shape[0] > shape[1]:
        ridge = Ridge(alpha=alpha, fit_intercept=fit_intercept)
        expected = ridge.fit(X, Y)
        kernel, bias = _compute_kernel_bias(X, Y, alpha=alpha, fit_intercept=fit_intercept)
        onp.testing.assert_array_almost_equal(expected.coef_, kernel, decimal=2)
        onp.testing.assert_array_almost_equal(expected.intercept_, bias, decimal=2)


@pytest.mark.parametrize("shape", [(40, 1), (50, 2)])
@pytest.mark.parametrize("alpha", [0.0, 0.5, 1.0])
@pytest.mark.parametrize("history_len", [1, 5, 10])
def test_ar_fit(shape, history_len, alpha):
    """Test AR class method fit"""
    X = jax.random.uniform(random.generate_key(), shape=shape)
    Y = jax.random.uniform(random.generate_key(), shape=(shape[0],))

    num_histories = shape[0] - history_len + 1
    history = historify(X, num_histories, history_len).reshape(num_histories, -1)

    kernel, bias = _compute_kernel_bias(
        history, Y[-len(history) :], alpha=alpha, fit_intercept=True
    )
    kernel = kernel.reshape(1, history_len * shape[1], 1)
    bias = jnp.expand_dims(jnp.asarray(bias), 0)

    ar, state = AR.fit(
        [(X, Y)], input_dim=X.shape[1], history_len=history_len, alpha=alpha, normalize=False
    )
    ar_kernel = ar.params["linear"]["kernel"]
    ar_bias = ar.params["linear"]["bias"]

    onp.testing.assert_array_almost_equal(kernel, ar_kernel, decimal=3)
    onp.testing.assert_array_almost_equal(bias, ar_bias, decimal=3)


def test_ar_fit_value_error():
    """Test number of observations"""
    X = jax.random.uniform(random.generate_key(), shape=(1, 10))
    Y = jax.random.uniform(random.generate_key(), shape=(10, 1))

    with pytest.raises(ValueError):
        AR.fit([(X, Y)], input_dim=10, history_len=1)


@pytest.mark.parametrize(
    "shape", [(1,), (1, 1), (10,), (10, 1), (1, 10), (10, 10), (20, 5), (5, 20)]
)
@pytest.mark.parametrize("history_len", [1, 5, 10])
@pytest.mark.parametrize("loc", [True, False])
@pytest.mark.parametrize("scale", [True, False])
def test_ar_apply(shape, history_len, loc, scale):
    """Test apply"""
    x = jax.random.uniform(random.generate_key(), shape=shape)
    if x.ndim == 1:
        x = x.reshape(1, -1)
    history = jax.random.uniform(random.generate_key(), shape=(history_len, x.shape[1]))
    kernel = jax.random.uniform(random.generate_key(), shape=(1, history_len * x.shape[1], 1))
    bias = jax.random.uniform(random.generate_key(), shape=(1,))

    loc = (
        jax.random.uniform(random.generate_key(), shape=(1, history_len * x.shape[1]))
        if loc
        else None
    )
    scale = (
        jax.random.uniform(random.generate_key(), shape=(1, history_len * x.shape[1]))
        if scale
        else None
    )

    with flax.nn.stateful() as _:
        model_def = AR.partial(
            history_len=history_len, history=history, loc=loc, scale=scale, name="AR"
        )
        _, params = model_def.init_by_shape(jax.random.PRNGKey(0), [shape])
        model = flax.nn.Model(model_def, params)
        model.params["linear"]["kernel"] = kernel
        model.params["linear"]["bias"] = bias

        result = model(x)

    expected_history = jnp.vstack((history, x))[x.shape[0] :].ravel()
    if loc is not None:
        expected_history -= loc
    if scale is not None:
        expected_history /= scale

    expected = jnp.tensordot(expected_history.reshape(1, -1), kernel) + bias

    onp.testing.assert_array_almost_equal(result, expected)
