"""
Testing configurations for the NeMoS library.

This module contains test fixtures required to set up and verify the functionality
of the modules of the NeMoS library.

Note:
    This module primarily serves as a utility for test configurations, setting up initial conditions,
    and loading predefined parameters for testing various functionalities of the NeMoS library.
"""

import abc
import os
from collections import namedtuple
from copy import deepcopy
from functools import partial
from typing import Literal

from nemos.glm.validation import GLMValidator

# Named tuple for model fixture returns (clearer than tuple indexing)
ModelFixture = namedtuple(
    "ModelFixture",
    ["X", "y", "model", "params", "rates", "extra"],
    defaults=[None, None],  # rates and extra default to None
)

import jax
import jax.numpy as jnp
import numpy as np
import pynapple as nap
import pytest

import nemos as nmo
import nemos._inspect_utils as inspect_utils
import nemos.basis.basis as basis
from nemos.basis import AdditiveBasis, CustomBasis, MultiplicativeBasis, Zero
from nemos.basis._basis import Basis
from nemos.basis._basis_mixin import BasisMixin
from nemos.basis._transformer_basis import TransformerBasis
from nemos.glm.params import GLMParams
from nemos.inverse_link_function_utils import log_softmax
from nemos.pytrees import FeaturePytree


def initialize_feature_mask_for_population_glm(X, n_neurons: int, coef=None):
    """
    Create a feature mask of ones for PopulationGLM testing.

    This is a test utility function that creates a feature mask with all ones,
    matching the structure of X.

    Parameters
    ----------
    X :
        The design matrix. Can be a FeaturePytree, dict, or array.
    n_neurons :
        Number of neurons (determines the second dimension of the mask).
        Ignored if coef is provided.
    coef :
        Optional coefficient array/pytree. If provided, the mask shape will match
        coef shape exactly (required for ClassifierPopulationGLM).

    Returns
    -------
    :
        A feature mask with all ones. If coef is provided, returns ones_like(coef).
        Otherwise, if X is a FeaturePytree or dict, returns a dict with arrays
        of shape (n_neurons,) for each key.
        If X is an array, returns an array of shape (n_features, n_neurons).
    """
    if coef is not None:
        return jax.tree_util.tree_map(lambda c: jnp.ones(c.shape), coef)
    if isinstance(X, FeaturePytree):
        return jax.tree_util.tree_map(lambda x: jnp.ones((n_neurons,)), X.data)
    elif isinstance(X, dict):
        return jax.tree_util.tree_map(lambda x: jnp.ones((n_neurons,)), X)
    else:
        return jnp.ones((X.shape[1], n_neurons))


DEFAULT_KWARGS = {
    "n_basis_funcs": 5,
    "frequencies": 4,
    "window_size": 11,
    "decay_rates": np.arange(1, 1 + 5),
}

# shut-off conversion warnings
nap.nap_config.suppress_conversion_warnings = True


@pytest.fixture(autouse=True, scope="function")
def set_jax_precision_per_test(request):
    """
    Automatically set JAX precision based on test marker.

    Tests marked with @pytest.mark.requires_x64 get float64.
    All other tests get float32 (JAX default).

    This fixture runs automatically for every test to ensure consistent
    precision behavior, especially important for parallel test execution
    with pytest-xdist.
    """
    if request.node.get_closest_marker("requires_x64"):
        # This test needs x64
        original = jax.config.jax_enable_x64
        jax.config.update("jax_enable_x64", True)
        yield
        jax.config.update("jax_enable_x64", original)
    else:
        # Default: float32
        original = jax.config.jax_enable_x64
        jax.config.update("jax_enable_x64", False)
        yield
        jax.config.update("jax_enable_x64", original)


@pytest.fixture()
def basis_class_specific_params():
    """Returns all the params for each class."""
    all_cls = (
        list_all_basis_classes("Conv")
        + list_all_basis_classes("Eval")
        + [CustomBasis, Zero]
    )
    return {cls.__name__: cls._get_param_names() for cls in all_cls}


class BasisFuncsTesting(abc.ABC):
    """
    An abstract base class that sets the foundation for individual basis function testing.
    This class requires an implementation of a 'cls' method, which is utilized by the meta-test
    that verifies if all basis functions are properly tested.
    """

    @abc.abstractmethod
    def cls(self):
        pass


def power_func(n, x):
    return jnp.power(x, n)


def custom_basis(n_basis_funcs=5, label=None, **kwargs):
    funcs = [partial(power_func, n) for n in range(1, n_basis_funcs + 1)]
    ndim_input = kwargs.get("ndim_input", 1)
    out_shape = kwargs.get("output_shape", None)
    pynapple_support = kwargs.get("pynapple_support", True)
    return CustomBasis(
        funcs,
        label=label,
        ndim_input=ndim_input,
        output_shape=out_shape,
        pynapple_support=pynapple_support,
    )


def power_add(n, x, y):
    return x**n + y**n


def custom_basis_2d(n_basis_funcs=5, label=None):
    funcs = [partial(power_add, n) for n in range(1, n_basis_funcs + 1)]
    return CustomBasis(funcs, label=label)


def basis_collapse_all_non_vec_axis(n_basis_funcs=5, label=None, **kwargs):
    funcs = [lambda x: x.reshape(x.shape[0], -1)[:, 0] for _ in range(n_basis_funcs)]
    ndim_input = kwargs.get("ndim_input", 1)
    return CustomBasis(funcs, label=label, ndim_input=ndim_input)


def basis_with_add_kwargs(label=None, basis_kwargs=None):
    def func(x, add=0):
        return x + add

    return CustomBasis([func], label=label, basis_kwargs=basis_kwargs)


class CombinedBasis(BasisFuncsTesting):
    """
    This class is used to run tests on combination operations (e.g., addition, multiplication) among Basis functions.

    Properties:
    - cls: Class (default = None)
    """

    cls = None

    @staticmethod
    def instantiate_basis(
        n_basis, basis_class, class_specific_params, window_size=10, **kwargs
    ):
        """Instantiate and return two basis of the type specified."""

        # Set non-optional args
        new_kwargs = {
            "n_basis_funcs": n_basis,
            "window_size": window_size,
            "decay_rates": np.arange(1, 1 + n_basis),
            "frequencies": np.arange(
                (n_basis + 1) % 2, 1 + (n_basis - n_basis % 2) // 2
            ),
            "frequency_mask": None,
        }
        default_kwargs = DEFAULT_KWARGS.copy()
        default_kwargs.update(new_kwargs)
        repeated_keys = set(new_kwargs.keys()).intersection(kwargs.keys())
        if repeated_keys:
            raise ValueError(
                "Cannot set `n_basis_funcs, window_size, decay_rates` with kwargs"
            )

        # Merge with provided  extra kwargs
        kwargs = {**default_kwargs, **kwargs}

        if basis_class == AdditiveBasis:
            kwargs_mspline = inspect_utils.trim_kwargs(
                basis.MSplineEval, kwargs, class_specific_params
            )
            kwargs_raised_cosine = inspect_utils.trim_kwargs(
                basis.RaisedCosineLinearConv, kwargs, class_specific_params
            )
            b1 = basis.MSplineEval(**kwargs_mspline)
            b2 = basis.RaisedCosineLinearConv(**kwargs_raised_cosine)
            basis_obj = b1 + b2
        elif basis_class == MultiplicativeBasis:
            kwargs_mspline = inspect_utils.trim_kwargs(
                basis.MSplineEval, kwargs, class_specific_params
            )
            kwargs_raised_cosine = inspect_utils.trim_kwargs(
                basis.RaisedCosineLinearConv, kwargs, class_specific_params
            )
            b1 = basis.MSplineEval(**kwargs_mspline)
            b2 = basis.RaisedCosineLinearConv(**kwargs_raised_cosine)
            basis_obj = b1 * b2
        elif basis_class == CustomBasis:
            basis_obj = custom_basis(
                n_basis,
                **inspect_utils.trim_kwargs(basis_class, kwargs, class_specific_params),
            )
        else:
            basis_obj = basis_class(
                **inspect_utils.trim_kwargs(basis_class, kwargs, class_specific_params)
            )
        return basis_obj


def is_eval_basis(basis_cls) -> bool:
    is_eval = "Eval" in basis_cls.__name__ or issubclass(basis_cls, basis.Zero)
    return is_eval


def is_conv_basis(basis_cls) -> bool:
    is_eval = "Conv" in basis_cls.__name__
    return is_eval


# automatic define user accessible basis and check the methods
def list_all_basis_classes(filter_basis="all") -> list[BasisMixin]:
    """
    Return all the classes in nemos.basis which are a subclass of Basis,
    which should be all concrete classes except TransformerBasis.
    """
    all_basis = (
        [
            class_obj
            for _, class_obj in inspect_utils.get_non_abstract_classes(basis)
            if issubclass(class_obj, Basis)
        ]
        + [
            bas
            for _, bas in inspect_utils.get_non_abstract_classes(nmo.basis._basis)
            if bas != TransformerBasis
        ]
        + [CustomBasis]
    )
    if filter_basis != "all":
        cond_fn = is_eval_basis if filter_basis == "Eval" else is_conv_basis
        all_basis = [a for a in all_basis if cond_fn(a)]
    return all_basis


def list_all_real_basis_classes(filter_basis="all"):
    list_all_basis = list_all_basis_classes(filter_basis)
    return [cls for cls in list_all_basis if not getattr(cls, "_is_complex", False)]


# Sample subclass to test instantiation and methods
class MockRegressor(nmo.base_regressor.BaseRegressor):
    """
    Mock implementation of the BaseRegressor abstract class for testing purposes.
    Implements all required abstract methods as empty methods.
    """

    _validator = GLMValidator()

    def __init__(self, std_param: int = 0):
        """Initialize a MockBaseRegressor instance with optional standard parameters."""
        self.std_param = std_param
        super().__init__()

    def fit(self, X, y):
        pass

    def predict(self, X) -> jnp.ndarray:
        pass

    def score(
        self,
        X,
        y,
        **kwargs,
    ) -> jnp.ndarray:
        pass

    def simulate(
        self,
        random_key: jax.Array,
        feed_forward_input,
        **kwargs,
    ):
        pass

    def _get_model_params(self):
        pass

    def _set_model_params(self, params):
        pass

    def update(self, *args, **kwargs):
        pass

    def _initialize_solver_and_state(self, *args, **kwargs):
        pass

    def initialize_params(self, *args, **kwargs):
        pass

    def _initialize_parameters(self, *args, **kwargs):
        pass

    def _compute_loss(self, params, X, y):
        pass

    def _get_optimal_solver_params_config(self):
        return None, None, None

    def save_params(self, *args):
        pass

    def _model_specific_initialization(self, *args, **kwargs):
        pass


class MockRegressorNested(MockRegressor):
    def __init__(self, other_param: int, std_param: int = 0):
        super().__init__(std_param=std_param)
        self.other_param = MockGLM(std_param=other_param)


class MockGLM(nmo.glm.GLM):
    """
    Mock implementation of the BaseRegressor abstract class for testing purposes.
    Implements all required abstract methods as empty methods.
    """

    def __init__(self, std_param: int = 0):
        """Initialize a MockBaseRegressor instance with optional standard parameters."""
        self.std_param = std_param
        super().__init__()

    def fit(self, X, y):
        pass

    def predict(self, X) -> jnp.ndarray:
        pass

    def score(
        self,
        X,
        y,
        **kwargs,
    ) -> jnp.ndarray:
        pass

    def simulate(
        self,
        random_key: jax.Array,
        feedforward_input,
        **kwargs,
    ):
        pass

    def _get_model_params(self):
        pass

    def _set_model_params(self, params):
        pass


@pytest.fixture
def mock_regressor():
    return MockRegressor(std_param=2)


@pytest.fixture
def mock_regressor_nested():
    return MockRegressorNested(other_param=1, std_param=2)


@pytest.fixture
def mock_glm():
    return MockGLM(std_param=2)


@pytest.fixture
def poissonGLM_model_instantiation():
    """Set up a Poisson GLM for testing purposes.

    This fixture initializes a Poisson GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (numpy.ndarray): Simulated input data.
            - np.random.poisson(rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.PoissonGLM): Initialized model instance.
            - GLMParams(w_true, b_true) (tuple): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of response.
    """
    np.random.seed(123)
    X = np.random.normal(size=(100, 5))
    b_true = np.zeros((1,))
    w_true = np.random.normal(size=(5,))
    observation_model = nmo.observation_models.PoissonObservations()
    regularizer = nmo.regularizer.UnRegularized()
    model = nmo.glm.GLM(observation_model, regularizer=regularizer)
    rate = jax.numpy.exp(jax.numpy.einsum("k,tk->t", w_true, X) + b_true)
    return X, np.random.poisson(rate), model, GLMParams(w_true, b_true), rate


@pytest.fixture
def poissonGLM_model_instantiation_pytree(poissonGLM_model_instantiation):
    """Set up a Poisson GLM for testing purposes.

    This fixture initializes a Poisson GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (numpy.ndarray): Simulated input data.
            - np.random.poisson(rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.PoissonGLM): Initialized model instance.
            - GLMParams(w_true, b_true) (tuple): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of response.
    """
    X, spikes, model, true_params, rate = poissonGLM_model_instantiation
    X_tree = nmo.pytrees.FeaturePytree(input_1=X[..., :3], input_2=X[..., 3:])
    true_params_tree = GLMParams(
        dict(input_1=true_params.coef[:3], input_2=true_params.coef[3:]),
        true_params.intercept,
    )
    model_tree = nmo.glm.GLM(model.observation_model, regularizer=model.regularizer)
    return X_tree, np.random.poisson(rate), model_tree, true_params_tree, rate


@pytest.fixture
def poissonGLM_fitted_model_instantiation(poissonGLM_model_instantiation):
    X, y, model, true_params, firing_rate = poissonGLM_model_instantiation
    model.fit(X, y)

    return X, y, model, true_params, firing_rate


@pytest.fixture
def population_poissonGLM_model_instantiation():
    """Set up a population Poisson GLM for testing purposes.

    This fixture initializes a Poisson GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (numpy.ndarray): Simulated input data.
            - np.random.poisson(rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.PoissonGLM): Initialized model instance.
            - GLMParams(w_true, b_true) (tuple): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of response.
    """
    np.random.seed(123)
    X = np.random.normal(size=(500, 5))
    b_true = -2 * np.ones((3,))
    w_true = np.random.normal(size=(5, 3))
    observation_model = nmo.observation_models.PoissonObservations()
    regularizer = nmo.regularizer.UnRegularized()
    model = nmo.glm.PopulationGLM(
        observation_model=observation_model, regularizer=regularizer
    )
    rate = jnp.exp(jnp.einsum("ki,tk->ti", w_true, X) + b_true)
    return X, np.random.poisson(rate), model, GLMParams(w_true, b_true), rate


@pytest.fixture
def population_poissonGLM_model_instantiation_pytree(
    population_poissonGLM_model_instantiation,
):
    """Set up a population Poisson GLM for testing purposes.

    This fixture initializes a Poisson GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (numpy.ndarray): Simulated input data.
            - np.random.poisson(rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.PoissonGLM): Initialized model instance.
            - GLMParams(w_true, b_true) (tuple): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of response.
    """
    X, spikes, model, true_params, rate = population_poissonGLM_model_instantiation
    X_tree = nmo.pytrees.FeaturePytree(input_1=X[..., :3], input_2=X[..., 3:])
    true_params_tree = GLMParams(
        dict(input_1=true_params.coef[:3], input_2=true_params.coef[3:]),
        true_params.intercept,
    )
    model_tree = nmo.glm.PopulationGLM(
        observation_model=model.observation_model, regularizer=model.regularizer
    )
    return X_tree, np.random.poisson(rate), model_tree, true_params_tree, rate


@pytest.fixture
def population_poissonGLM_fitted_model_instantiation(
    population_poissonGLM_model_instantiation,
):
    X, y, model, true_params, firing_rate = population_poissonGLM_model_instantiation
    model.fit(X, y)

    return X, y, model, true_params, firing_rate


@pytest.fixture
def coupled_model_simulate():
    """Set up a Poisson GLM from a predefined configuration in a json file.

    This fixture reads parameters for a Poisson GLM from a json configuration file, initializes
    the model accordingly, and returns the model instance with other related parameters.

    Returns:
        tuple: A tuple containing:
            - the coupling coeffs
            - the feedforward coeffs
            - the intercepts
            - jax.random.key(123) (jax.Array): A pseudo-random number generator key.
            - feedforward_input (jax.numpy.ndarray): Feedforward input values from the config.
            - coupling_basis (jax.numpy.ndarray): Coupling basis values from the config.
            - init_spikes (jax.numpy.ndarray): Initial spike values from the config.
            - a link function.

    """

    n_neurons, coupling_duration, sim_duration = 2, 100, 1000
    coupling_filter_bank = np.zeros((coupling_duration, n_neurons, n_neurons))
    for unit_i in range(n_neurons):
        for unit_j in range(n_neurons):
            coupling_filter_bank[:, unit_i, unit_j] = (
                nmo.simulation.difference_of_gammas(coupling_duration)
            )
    # shrink the filters for simulation stability
    coupling_filter_bank *= 0.8
    basis = nmo.basis.RaisedCosineLogEval(20)

    # approximate the coupling filters in terms of the basis function
    _, coupling_basis = basis.evaluate_on_grid(coupling_filter_bank.shape[0])
    coupling_coeff = nmo.simulation.regress_filter(coupling_filter_bank, coupling_basis)
    feedforward_coeff = np.ones((n_neurons, 2))
    intercepts = -3 * jnp.ones(n_neurons)

    feedforward_input = jnp.c_[
        jnp.cos(jnp.linspace(0, np.pi * 4, sim_duration)),
        jnp.sin(jnp.linspace(0, np.pi * 4, sim_duration)),
    ]

    feedforward_input = jnp.tile(feedforward_input[:, None], (1, n_neurons, 1))
    init_spikes = jnp.zeros((coupling_duration, n_neurons))

    return (
        coupling_coeff,
        feedforward_coeff,
        intercepts,
        jax.random.key(123),
        feedforward_input,
        coupling_basis,
        init_spikes,
        jnp.exp,
    )


@pytest.fixture
def poissonGLM_model_instantiation_group_sparse():
    """Set up a Poisson GLM for testing purposes with group sparse weights.

    This fixture initializes a Poisson GLM with random, group sparse, parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response

    Returns:
        tuple: A tuple containing:
            - X (numpy.ndarray): Simulated input data.
            - np.random.poisson(rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.PoissonGLM): Initialized model instance.
            - GLMParams(w_true, b_true) (tuple): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of response.
    """
    np.random.seed(123)
    X = np.random.normal(size=(100, 5))
    b_true = np.zeros((1,))
    w_true = np.random.normal(size=(5,))
    w_true[1:4] = 0.0
    mask = np.zeros((2, 5))
    mask[0, 1:4] = 1
    mask[1, [0, 4]] = 1
    observation_model = nmo.observation_models.PoissonObservations()
    regularizer = nmo.regularizer.UnRegularized()
    model = nmo.glm.GLM(observation_model, regularizer=regularizer)
    rate = jax.numpy.exp(jax.numpy.einsum("k,tk->t", w_true, X) + b_true)
    return X, np.random.poisson(rate), model, GLMParams(w_true, b_true), rate, mask


@pytest.fixture
def population_poissonGLM_model_instantiation_group_sparse():
    """Set up a Poisson GLM for testing purposes with group sparse weights.

    This fixture initializes a Poisson GLM with random, group sparse, parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response

    Returns:
        tuple: A tuple containing:
            - X (numpy.ndarray): Simulated input data.
            - np.random.poisson(rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.PoissonGLM): Initialized model instance.
            - GLMParams(w_true, b_true) (tuple): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of response.
    """
    np.random.seed(123)
    X = np.random.normal(size=(100, 5))
    b_true = np.zeros((3,))
    w_true = np.random.normal(size=(5, 3))
    w_true[1:4, 0] = 0.0
    mask = np.zeros((2, 5))
    mask[0, 1:4] = 1
    mask[1, [0, 4]] = 1
    observation_model = nmo.observation_models.PoissonObservations()
    regularizer = nmo.regularizer.UnRegularized()
    model = nmo.glm.PopulationGLM(observation_model, regularizer=regularizer)
    rate = jax.numpy.exp(jax.numpy.einsum("kn,tk->tn", w_true, X) + b_true)
    return X, np.random.poisson(rate), model, GLMParams(w_true, b_true), rate, mask


@pytest.fixture
def example_data_prox_operator():
    n_features = 4

    params = GLMParams(
        jnp.ones((n_features)),
        jnp.zeros(1),
    )
    regularizer_strength = 0.1
    # Mask as PyTree with same structure as params, shape (n_groups, *param_shape)
    # Intercept mask is zeros (not regularized)
    mask = GLMParams(
        jnp.array([[1, 0, 1, 0], [0, 1, 0, 1]]).astype(float),
        jnp.zeros((2, 1), dtype=float),
    )
    scaling = 0.5

    return params, regularizer_strength, mask, scaling


@pytest.fixture
def example_data_prox_operator_multineuron():
    n_features = 4
    n_neurons = 3

    params = GLMParams(
        jnp.ones((n_features, n_neurons)),
        jnp.zeros(n_neurons),
    )
    regularizer_strength = 0.1
    # Mask as PyTree with same structure as params
    # For multi-neuron: mask shape is (n_groups, n_features, n_neurons)
    # Intercept mask is zeros (not regularized)
    mask_coef = jnp.array(
        [
            [[1, 1, 1], [0, 0, 0], [1, 1, 1], [0, 0, 0]],
            [[0, 0, 0], [1, 1, 1], [0, 0, 0], [1, 1, 1]],
        ],
        dtype=jnp.float32,
    )
    mask = GLMParams(
        mask_coef,
        jnp.zeros((2, n_neurons), dtype=jnp.float32),
    )
    scaling = 0.5

    return params, regularizer_strength, mask, scaling


@pytest.fixture
def poisson_observation_model():
    return nmo.observation_models.PoissonObservations(jnp.exp)


@pytest.fixture
def ridge_regularizer():
    return nmo.regularizer.Ridge()


@pytest.fixture
def lasso_regularizer():
    return nmo.regularizer.Lasso(solver_name="ProximalGradient")


@pytest.fixture
def group_lasso_2groups_5features_regularizer():
    mask = np.zeros((2, 5))
    mask[0, :2] = 1
    mask[1, 2:] = 1
    return nmo.regularizer.GroupLasso(solver_name="ProximalGradient", mask=mask)


@pytest.fixture
def mock_data():
    return jnp.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]), jnp.array([[1, 2], [3, 4]])


@pytest.fixture()
def glm_class():
    return nmo.glm.GLM


@pytest.fixture()
def population_glm_class():
    return nmo.glm.PopulationGLM


@pytest.fixture
def gammaGLM_model_instantiation():
    """Set up a Gamma GLM for testing purposes.

    This fixture initializes a Gamma GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (numpy.ndarray): Simulated input data.
            - np.random.poisson(rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.PoissonGLM): Initialized model instance.
            - GLMParams(w_true, b_true) (tuple): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of response.
    """
    np.random.seed(123)
    X = np.random.uniform(size=(100, 5))
    b_true = np.zeros((1,))
    w_true = np.random.uniform(size=(5,))
    observation_model = nmo.observation_models.GammaObservations()
    regularizer = nmo.regularizer.UnRegularized()
    model = nmo.glm.GLM(observation_model, regularizer=regularizer)
    rate = (jax.numpy.einsum("k,tk->t", w_true, X) + b_true) ** -1
    theta = 3
    k = rate / theta
    model.scale_ = theta
    return X, np.random.gamma(k, scale=theta), model, GLMParams(w_true, b_true), rate


@pytest.fixture
def gammaGLM_model_instantiation_pytree(gammaGLM_model_instantiation):
    """Set up a Gamma GLM for testing purposes.

    This fixture initializes a Gamma GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (numpy.ndarray): Simulated input data.
            - np.random.poisson(rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.PoissonGLM): Initialized model instance.
            - GLMParams(w_true, b_true) (tuple): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of response.
    """
    X, spikes, model, true_params, rate = gammaGLM_model_instantiation
    X_tree = nmo.pytrees.FeaturePytree(input_1=X[..., :3], input_2=X[..., 3:])
    true_params_tree = GLMParams(
        dict(input_1=true_params.coef[:3], input_2=true_params.coef[3:]),
        true_params.intercept,
    )
    model_tree = nmo.glm.GLM(model.observation_model, regularizer=model.regularizer)
    return X_tree, spikes, model_tree, true_params_tree, rate


@pytest.fixture()
def population_gammaGLM_model_instantiation():
    np.random.seed(123)
    X = np.random.uniform(size=(500, 5))
    b_true = 0.5 * np.ones((3,))
    w_true = np.random.uniform(size=(5, 3))
    observation_model = nmo.observation_models.GammaObservations()
    regularizer = nmo.regularizer.UnRegularized()
    model = nmo.glm.PopulationGLM(
        observation_model=observation_model, regularizer=regularizer
    )
    rate = 1 / (jnp.einsum("ki,tk->ti", w_true, X) + b_true)
    theta = 3
    model.scale_ = theta
    y = jax.random.gamma(jax.random.PRNGKey(123), rate / theta) * theta
    return X, y, model, GLMParams(w_true, b_true), rate


@pytest.fixture()
def population_gammaGLM_model_instantiation_pytree(
    population_gammaGLM_model_instantiation,
):
    X, spikes, model, true_params, rate = population_gammaGLM_model_instantiation
    X_tree = nmo.pytrees.FeaturePytree(input_1=X[..., :3], input_2=X[..., 3:])
    true_params_tree = GLMParams(
        dict(input_1=true_params.coef[:3], input_2=true_params.coef[3:]),
        true_params.intercept,
    )
    model_tree = nmo.glm.PopulationGLM(
        observation_model=model.observation_model, regularizer=model.regularizer
    )
    return X_tree, spikes, model_tree, true_params_tree, rate


@pytest.fixture
def regr_data():
    np.random.seed(123)
    # define inputs and coeff
    n_samples, n_features = 50, 3
    X = np.random.normal(size=(n_samples, n_features))
    coef = np.random.normal(size=(n_features))
    # set y according to lin reg eqn
    y = X.dot(coef) + 0.1 * np.random.normal(size=(n_samples,))
    return X, y, coef


@pytest.fixture
def linear_regression(regr_data):
    X, y, coef = regr_data
    # solve least-squares
    ols, _, _, _ = np.linalg.lstsq(X, y, rcond=-1)

    # set the loss
    def loss(params, X, y):
        return jnp.power(y - jnp.dot(X, params), 2).mean()

    return X, y, coef, ols, loss


@pytest.fixture
def ridge_regression(regr_data):
    X, y, coef = regr_data

    # solve least-squares
    yagu = np.hstack((y, np.zeros_like(coef)))
    Xagu = np.vstack((X, np.sqrt(0.5) * np.eye(coef.shape[0])))
    ridge, _, _, _ = np.linalg.lstsq(Xagu, yagu, rcond=-1)

    # set the loss
    def loss(params, XX, yy):
        return (
            jnp.power(yy - jnp.dot(XX, params), 2).sum()
            + 0.5 * jnp.power(params, 2).sum()
        )

    return X, y, coef, ridge, loss


@pytest.fixture
def linear_regression_tree(linear_regression):
    X, y, coef, ols, loss = linear_regression
    X_tree = dict(input_1=X[..., :2], input_2=X[..., 2:])
    coef_tree = dict(input_1=coef[:2], input_2=coef[2:])
    ols_tree = dict(input_1=ols[:2], input_2=ols[2:])

    nmo.tree_utils.pytree_map_and_reduce(jnp.dot, sum, X_tree, coef_tree)

    def loss_tree(params, XX, yy):
        pred = nmo.tree_utils.pytree_map_and_reduce(jnp.dot, sum, XX, params)
        return jnp.power(yy - pred, 2).sum()

    return X_tree, y, coef_tree, ols_tree, loss_tree


@pytest.fixture()
def ridge_regression_tree(ridge_regression):
    X, y, coef, ridge, loss = ridge_regression
    X_tree = dict(input_1=X[..., :2], input_2=X[..., 2:])
    coef_tree = dict(input_1=coef[:2], input_2=coef[2:])
    ridge_tree = dict(input_1=ridge[:2], input_2=ridge[2:])

    def loss_tree(params, XX, yy):
        pred = nmo.tree_utils.pytree_map_and_reduce(jnp.dot, sum, XX, params)
        norm = (
            0.5
            * nmo.tree_utils.pytree_map_and_reduce(
                lambda x: jnp.power(x, 2).sum(), sum, params
            ).sum()
        )
        return jnp.power(yy - pred, 2).sum() + norm

    return X_tree, y, coef_tree, ridge_tree, loss_tree


@pytest.fixture()
def example_X_y_high_firing_rates():
    """Example that used failed with NeMoS original initialization."""
    np.random.seed(123)

    n_features = 18
    n_neurons = 60
    n_samples = 500

    # random design array. Shape (n_time_points, n_features).
    X = 0.5 * np.random.normal(size=(n_samples, n_features))

    # log-rates & weights
    b_true = np.random.uniform(size=(n_neurons,)) * 3  # baseline rates
    w_true = np.random.uniform(size=(n_features, n_neurons)) * 0.1  # real weights:

    # generate counts (spikes will be (n_samples, n_features)
    rate = jnp.exp(jnp.dot(X, w_true) + b_true)
    spikes = np.random.poisson(rate)
    return X, spikes


@pytest.fixture
def bernoulliGLM_model_instantiation():
    """Set up a Bernoulli GLM for testing purposes.

    This fixture initializes a Bernoulli GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (numpy.ndarray): Simulated input data.
            - np.random.binomial(1,rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.GLM): Initialized model instance.
            - GLMParams(w_true, b_true) (tuple): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of response.
    """
    np.random.seed(123)
    X = np.random.normal(size=(100, 5))
    b_true = np.zeros((1,))
    w_true = np.random.normal(size=(5,))
    observation_model = nmo.observation_models.BernoulliObservations()
    regularizer = nmo.regularizer.UnRegularized()
    model = nmo.glm.GLM(observation_model, regularizer=regularizer)
    rate = jax.lax.logistic(jnp.einsum("k,tk->t", w_true, X) + b_true)
    return X, np.random.binomial(1, rate), model, GLMParams(w_true, b_true), rate


@pytest.fixture
def bernoulliGLM_model_instantiation_pytree(bernoulliGLM_model_instantiation):
    """Set up a Bernoulli GLM for testing purposes.

    This fixture initializes a Bernoulli GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (numpy.ndarray): Simulated input data.
            - np.random.binomial(1,rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.GLM): Initialized model instance.
            - GLMParams(w_true, b_true) (tuple): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of response.
    """
    X, spikes, model, true_params, rate = bernoulliGLM_model_instantiation
    X_tree = nmo.pytrees.FeaturePytree(input_1=X[..., :3], input_2=X[..., 3:])
    true_params_tree = GLMParams(
        dict(input_1=true_params.coef[:3], input_2=true_params.coef[3:]),
        true_params.intercept,
    )
    model_tree = nmo.glm.GLM(model.observation_model, regularizer=model.regularizer)
    return X_tree, spikes, model_tree, true_params_tree, rate


@pytest.fixture
def classifierGLM_model_instantiation():
    """Set up a categorical GLM for testing purposes.

    This fixture initializes a categorical GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (numpy.ndarray): Simulated input data.
            - jax.random.categorical(key, rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.GLM): Initialized model instance.
            - GLMParams(w_true, b_true) (tuple): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of log-proba.
    """
    np.random.seed(123)
    num_classes = 3
    X = np.random.normal(size=(100, 5))
    # Over-parameterized model: one set of params per class
    b_true = np.zeros((num_classes,))
    w_true = np.random.normal(size=(5, num_classes))
    regularizer = nmo.regularizer.UnRegularized()
    model = nmo.glm.ClassifierGLM(num_classes, regularizer=regularizer)
    # Use jax.nn.log_softmax directly (no padding)
    rate = jax.nn.log_softmax(jnp.einsum("ki,tk->ti", w_true, X) + b_true)
    key = jax.random.PRNGKey(123)
    y = jax.random.categorical(key, rate)
    model.set_classes(np.arange(num_classes))
    return X, y, model, GLMParams(w_true, b_true), rate


@pytest.fixture
def classifierGLM_model_instantiation_pytree(classifierGLM_model_instantiation):
    """Set up a categorical GLM for testing purposes.

    This fixture initializes a categorical GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (numpy.ndarray): Simulated input data.
            - jax.random.categorical(key, rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.GLM): Initialized model instance.
            - GLMParams(w_true, b_true) (tuple): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of log-proba.
    """
    X, spikes, model, true_params, rate = classifierGLM_model_instantiation
    X_tree = nmo.pytrees.FeaturePytree(input_1=X[..., :3], input_2=X[..., 3:])
    true_params_tree = GLMParams(
        dict(input_1=true_params.coef[:3], input_2=true_params.coef[3:]),
        true_params.intercept,
    )
    model_tree = nmo.glm.ClassifierGLM(
        n_classes=model.n_classes, regularizer=model.regularizer
    )
    model_tree.set_classes(np.arange(model.n_classes))
    return X_tree, spikes, model_tree, true_params_tree, rate


@pytest.fixture
def population_classifierGLM_model_instantiation():
    """Set up a categorical GLM for testing purposes.

    This fixture initializes a categorical GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (numpy.ndarray): Simulated input data.
            - jax.random.categorical(key, rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.GLM): Initialized model instance.
            - GLMParams(w_true, b_true) (tuple): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of log-proba.
    """
    np.random.seed(123)
    num_classes = 3
    n_neurons = 3
    X = np.random.normal(size=(100, 5))
    # Over-parameterized model: one set of params per class
    b_true = np.zeros((n_neurons, num_classes))
    w_true = np.random.normal(size=(5, n_neurons, num_classes))
    regularizer = nmo.regularizer.UnRegularized()
    model = nmo.glm.ClassifierPopulationGLM(num_classes, regularizer=regularizer)
    # Use jax.nn.log_softmax directly (no padding)
    rate = jax.nn.log_softmax(jnp.einsum("kni,tk->tni", w_true, X) + b_true)
    key = jax.random.PRNGKey(123)
    y = jax.random.categorical(key, rate)
    model.set_classes(np.arange(num_classes))
    return X, y, model, GLMParams(w_true, b_true), rate


@pytest.fixture
def population_classifierGLM_model_instantiation_pytree(
    population_classifierGLM_model_instantiation,
):
    """Set up a categorical GLM for testing purposes.

    This fixture initializes a categorical GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (numpy.ndarray): Simulated input data.
            - jax.random.categorical(key, rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.GLM): Initialized model instance.
            - GLMParams(w_true, b_true) (tuple): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of log-proba.
    """
    X, spikes, model, true_params, rate = population_classifierGLM_model_instantiation
    X_tree = nmo.pytrees.FeaturePytree(input_1=X[..., :3], input_2=X[..., 3:])
    true_params_tree = GLMParams(
        dict(input_1=true_params.coef[:3], input_2=true_params.coef[3:]),
        true_params.intercept,
    )
    model_tree = nmo.glm.ClassifierPopulationGLM(
        n_classes=model.n_classes, regularizer=model.regularizer
    )
    model_tree.set_classes(np.arange(model.n_classes))
    return X_tree, spikes, model_tree, true_params_tree, rate


@pytest.fixture
def population_bernoulliGLM_model_instantiation():
    """Set up a population Bernoulli GLM for testing purposes.

    This fixture initializes a Bernoulli GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (numpy.ndarray): Simulated input data.
            - np.random.binomial(1,rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.PopulationGLM): Initialized model instance.
            - GLMParams(w_true, b_true) (tuple): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of response.
    """
    np.random.seed(123)
    X = np.random.normal(size=(500, 5))
    b_true = np.zeros((3,))
    w_true = np.random.normal(size=(5, 3))
    observation_model = nmo.observation_models.BernoulliObservations()
    regularizer = nmo.regularizer.UnRegularized()
    model = nmo.glm.PopulationGLM(
        observation_model=observation_model, regularizer=regularizer
    )
    rate = jax.lax.logistic(jnp.einsum("ki,tk->ti", w_true, X) + b_true)
    return X, np.random.binomial(1, rate), model, GLMParams(w_true, b_true), rate


@pytest.fixture
def population_bernoulliGLM_model_instantiation_pytree(
    population_bernoulliGLM_model_instantiation,
):
    """Set up a population Bernoulli GLM for testing purposes.

    This fixture initializes a Bernoulli GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (numpy.ndarray): Simulated input data.
            - np.random.binomial(1,rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.PopulationGLM): Initialized model instance.
            - GLMParams(w_true, b_true) (tuple): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of response.
    """
    X, spikes, model, true_params, rate = population_bernoulliGLM_model_instantiation
    X_tree = nmo.pytrees.FeaturePytree(input_1=X[..., :3], input_2=X[..., 3:])
    true_params_tree = GLMParams(
        dict(input_1=true_params.coef[:3], input_2=true_params.coef[3:]),
        true_params.intercept,
    )
    model_tree = nmo.glm.PopulationGLM(
        observation_model=model.observation_model, regularizer=model.regularizer
    )
    return X_tree, np.random.binomial(1, rate), model_tree, true_params_tree, rate


SizeTerminal = namedtuple("SizeTerminal", ["columns", "lines"])


class NestedRegularizer(nmo.regularizer.Ridge):
    def __init__(self, sub_regularizer, func=np.exp):
        self.sub_regularizer = sub_regularizer
        self.func = func
        super().__init__()


@pytest.fixture
def nested_regularizer():
    """
    Nested retularizer for testing save/load.
    """
    return NestedRegularizer(nmo.regularizer.Lasso())


@pytest.fixture
def negativeBinomialGLM_model_instantiation():
    """Set up a Negative Binomial GLM with array inputs for testing purposes.

    This fixture initializes a Negative Binomial GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (numpy.ndarray): Simulated input data.
            - np.random.poisson(rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.PoissonGLM): Initialized model instance.
            - GLMParams(w_true, b_true) (tuple): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of response.
    """
    np.random.seed(123)
    X = np.random.normal(size=(100, 5))
    b_true = np.zeros((1,))
    w_true = np.random.normal(size=(5,))
    observation_model = nmo.observation_models.NegativeBinomialObservations()
    regularizer = nmo.regularizer.UnRegularized()
    model = nmo.glm.GLM(observation_model, regularizer=regularizer, solver_name="LBFGS")
    rate = jax.numpy.exp(jax.numpy.einsum("k,tk->t", w_true, X) + b_true)
    r = 1 / model.observation_model.scale
    spikes = np.random.poisson(np.random.gamma(shape=r, size=rate.shape) * (r / rate))
    return X, spikes, model, GLMParams(w_true, b_true), rate


@pytest.fixture
def negativeBinomialGLM_model_instantiation_pytree(
    negativeBinomialGLM_model_instantiation,
):
    """Set up a Negative Binomial GLM with pytree inputs for testing purposes .

    This fixture initializes a Negative Binomial GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (FeaturePytree): Simulated input data.
            - np.random.poisson(rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.PoissonGLM): Initialized model instance.
            - GLMParams(w_true, b_true) (tuple): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of response.
    """
    X, spikes, model, true_params, rate = negativeBinomialGLM_model_instantiation
    X_tree = nmo.pytrees.FeaturePytree(input_1=X[..., :3], input_2=X[..., 3:])
    true_params_tree = GLMParams(
        dict(input_1=true_params.coef[:3], input_2=true_params.coef[3:]),
        true_params.intercept,
    )
    model_tree = nmo.glm.GLM(
        model.observation_model, regularizer=model.regularizer, solver_name="LBFGS"
    )
    return X_tree, np.random.poisson(rate), model_tree, true_params_tree, rate


@pytest.fixture
def population_negativeBinomialGLM_model_instantiation():
    """Set up a population Negative Binomial GLM for testing purposes.

    This fixture initializes a Negative Binomial GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (numpy.ndarray): Simulated input data.
            - np.random.poisson(rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.PoissonGLM): Initialized model instance.
            - GLMParams(w_true, b_true) (tuple): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of response.
    """
    np.random.seed(123)
    X = np.random.normal(size=(500, 5))
    b_true = -2 * np.ones((3,))
    w_true = 0.1 * np.random.normal(size=(5, 3))
    observation_model = nmo.observation_models.NegativeBinomialObservations()
    regularizer = nmo.regularizer.UnRegularized()
    model = nmo.glm.PopulationGLM(
        observation_model=observation_model,
        regularizer=regularizer,
        solver_name="LBFGS",
    )
    rate = jnp.exp(jnp.einsum("ki,tk->ti", w_true, X) + b_true)
    spikes = model.observation_model.sample_generator(jax.random.PRNGKey(123), rate)
    # make sure that at least one entry is non-zero
    spikes = spikes.at[-1].set(1)
    return X, spikes, model, GLMParams(w_true, b_true), rate


@pytest.fixture
def population_negativeBinomialGLM_model_instantiation_pytree(
    population_negativeBinomialGLM_model_instantiation,
):
    """Set up a population Negative Binomial GLM for testing purposes.

    This fixture initializes a Negative Binomial GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (numpy.ndarray): Simulated input data.
            - np.random.poisson(rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.PoissonGLM): Initialized model instance.
            - GLMParams(w_true, b_true) (tuple): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of response.
    """
    X, spikes, model, true_params, rate = (
        population_negativeBinomialGLM_model_instantiation
    )
    X_tree = nmo.pytrees.FeaturePytree(input_1=X[..., :3], input_2=X[..., 3:])
    true_params_tree = GLMParams(
        dict(input_1=true_params.coef[:3], input_2=true_params.coef[3:]),
        true_params.intercept,
    )
    model_tree = nmo.glm.PopulationGLM(
        observation_model=model.observation_model,
        regularizer=model.regularizer,
        solver_name="LBFGS",
    )
    return X_tree, np.random.poisson(rate), model_tree, true_params_tree, rate


def instantiate_glm_func(
    obs_model: (
        Literal["Poisson", "Gamma", "Bernoulli", "NegativeBinomial"]
        | nmo.observation_models.Observations
    ) = "Bernoulli",
    regularizer: str = "UnRegularized",
    solver_name: str = None,
    simulate=False,
):
    np.random.seed(123)
    n_features = 2
    X = np.ones((500, n_features))
    X[:250, 0] = 0
    X[np.arange(500) % 2 == 1, 1] = 0
    model = nmo.glm.GLM(
        observation_model=obs_model,
        regularizer=regularizer,
        solver_name=solver_name,
    )
    model.coef_ = np.random.randn(n_features)
    model.intercept_ = np.random.randn(1)
    if simulate:
        counts, rates = model.simulate(jax.random.PRNGKey(1234), X)
    else:
        counts, rates = None, None
    return ModelFixture(
        X=X,
        y=counts,
        model=model,
        params=GLMParams(model.coef_, model.intercept_),
        rates=rates,
        extra=None,
    )


def instantiate_population_glm_func(
    n_neurons=3,
    obs_model: (
        Literal["Poisson", "Gamma", "Bernoulli", "NegativeBinomial"]
        | nmo.observation_models.Observations
    ) = "Bernoulli",
    regularizer: str = "UnRegularized",
    solver_name: str = None,
    simulate=False,
):
    np.random.seed(123)
    n_features = 2
    X = np.ones((500, n_features))
    X[:250, 0] = 0
    X[np.arange(500) % 2 == 1, 1] = 0
    model = nmo.glm.PopulationGLM(
        observation_model=obs_model,
        regularizer=regularizer,
        solver_name=solver_name,
    )
    model.coef_ = np.random.randn(n_features, n_neurons)
    model.intercept_ = np.random.randn(n_neurons)
    if simulate:
        model._feature_mask = initialize_feature_mask_for_population_glm(X, n_neurons)
        counts, rates = model.simulate(jax.random.PRNGKey(1234), X)
    else:
        counts, rates = None, None
    return ModelFixture(
        X=X,
        y=counts,
        model=model,
        params=GLMParams(model.coef_, model.intercept_),
        rates=rates,
        extra=None,
    )


def instantiate_classifier_glm_func(
    n_neurons=3,
    regularizer: str = "UnRegularized",
    solver_name: str = None,
    simulate=False,
):
    np.random.seed(124)
    n_features = 2
    n_classes = 4
    X = np.ones((500, n_features))
    X[:250, 0] = 0
    X[np.arange(500) % 2 == 1, 1] = 0
    model = nmo.glm.ClassifierGLM(
        n_classes=n_classes,
        regularizer=regularizer,
        solver_name=solver_name,
    )
    model.coef_ = np.random.randn(n_features, n_classes)
    model.intercept_ = np.random.randn(n_classes)
    model.set_classes(np.arange(n_classes))
    if simulate:
        counts, rates = model.simulate(jax.random.PRNGKey(123), X)
    else:
        counts, rates = None, None
    return ModelFixture(
        X=X,
        y=counts,
        model=model,
        params=GLMParams(model.coef_, model.intercept_),
        rates=rates,
        extra=None,
    )


def instantiate_population_classifier_glm_func(
    n_neurons=3,
    regularizer: str = "UnRegularized",
    solver_name: str = None,
    simulate=False,
):
    np.random.seed(124)
    n_features = 2
    n_classes = 4
    X = np.ones((500, n_features))
    X[:250, 0] = 0
    X[np.arange(500) % 2 == 1, 1] = 0
    model = nmo.glm.ClassifierPopulationGLM(
        n_classes=n_classes,
        regularizer=regularizer,
        solver_name=solver_name,
    )
    model.set_classes(np.arange(n_classes))
    model.coef_ = np.random.randn(n_features, n_neurons, n_classes)
    model.intercept_ = np.random.randn(n_neurons, n_classes)
    if simulate:
        model._feature_mask = initialize_feature_mask_for_population_glm(X, n_neurons)
        counts, rates = model.simulate(jax.random.PRNGKey(123), X)
    else:
        counts, rates = None, None
    return ModelFixture(
        X=X,
        y=counts,
        model=model,
        params=GLMParams(model.coef_, model.intercept_),
        rates=rates,
        extra=None,
    )


_MODEL_CACHE = {}


# Registry for model-specific configurations
MODEL_CONFIG = {
    "GLM": {
        "is_population": False,
        "default_y_shape": (500,),
    },
    "ClassifierGLM": {
        "is_population": False,
        "default_y_shape": (500,),
    },
    "PopulationGLM": {
        "is_population": True,
        "default_y_shape": (500, 3),
    },
    "ClassifierPopulationGLM": {
        "is_population": True,
        "default_y_shape": (500, 3),
    },
}


def is_population_model(model) -> bool:
    """Check if a model is a population model using registry instead of string matching."""
    model_name = model.__class__.__name__
    return MODEL_CONFIG.get(model_name, {}).get("is_population", False)


@pytest.fixture
def instantiate_base_regressor_subclass(request):
    """
    Instantiate the concrete BaseRegressor sub-classes with caching.
    """
    model_name: str = request.param["model"]
    obs_model: str | nmo.observation_models.Observations = request.param["obs_model"]
    simulate: bool = request.param["simulate"]

    # Create cache key (class-scoped)
    cache_key = (
        model_name,
        str(obs_model),
        simulate,
        id(request.cls) if request.cls else id(request.module),
    )

    # Check cache
    if cache_key not in _MODEL_CACHE:
        if model_name == "GLM":
            result = instantiate_glm_func(obs_model=obs_model, simulate=simulate)
        elif model_name == "PopulationGLM":
            result = instantiate_population_glm_func(
                obs_model=obs_model, simulate=simulate
            )
        elif model_name == "ClassifierGLM":
            result = instantiate_classifier_glm_func(simulate=simulate)
        elif model_name == "ClassifierPopulationGLM":
            result = instantiate_population_classifier_glm_func(simulate=simulate)
        else:
            raise ValueError("model_name {} unknown".format(model_name))
        _MODEL_CACHE[cache_key] = result
        return deepcopy(result)

    # Get cached data and return a complete deepcopy of everything
    # this is different from a function level fixture because it
    # would not re-run any potentially heavy setup code (like model.simulate).
    cached_result = deepcopy(_MODEL_CACHE[cache_key])
    return cached_result


# Auto-clear cache after each test module run
@pytest.fixture(scope="module", autouse=True)
def _clear_model_cache():
    """Clear model cache after each test module."""
    yield
    _MODEL_CACHE.clear()


# Select solver backend for tests if requested via environment variable
_common_solvers = {
    "SVRG": nmo.solvers.WrappedSVRG,
    "ProxSVRG": nmo.solvers.WrappedProxSVRG,
}
_solver_registry_per_backend = {
    "optimistix": {
        **_common_solvers,
        "GradientDescent": nmo.solvers.OptimistixNAG,
        "ProximalGradient": nmo.solvers.OptimistixFISTA,
        "LBFGS": nmo.solvers.OptimistixOptaxLBFGS,
        "BFGS": nmo.solvers.OptimistixBFGS,
        "NonlinearCG": nmo.solvers.OptimistixNonlinearCG,
    },
}

if nmo.solvers.JAXOPT_AVAILABLE:
    _solver_registry_per_backend["jaxopt"] = {
        **_common_solvers,
        "GradientDescent": nmo.solvers.JaxoptGradientDescent,
        "ProximalGradient": nmo.solvers.JaxoptProximalGradient,
        "LBFGS": nmo.solvers.JaxoptLBFGS,
        "BFGS": nmo.solvers.JaxoptBFGS,
        "NonlinearCG": nmo.solvers.JaxoptNonlinearCG,
    }


@pytest.fixture(autouse=True, scope="session")
def configure_solver_backend(request):
    """
    Patch the solver registry depending on `NEMOS_SOLVER_BACKEND` and `override_solver`.

    The `NEMOS_SOLVER_BACKEND` env variable is used for running solver-dependent tests
    in separate tox environments for the JAXopt and the Optimistix backends.

    The `override_solver` pytest option is used to set a given solver algorithm's
    implementation to a class available in nemos.solvers.
    """
    backend = os.getenv("NEMOS_SOLVER_BACKEND")

    if backend is None:
        _solver_registry_to_use = nmo.solvers.solver_registry.copy()
    else:
        if backend == "jaxopt" and not nmo.solvers.JAXOPT_AVAILABLE:
            pytest.fail("jaxopt backend requested but jaxopt is not installed.")
        try:
            _solver_registry_to_use = _solver_registry_per_backend[backend]
        except KeyError:
            available = ", ".join(_solver_registry_per_backend.keys())
            pytest.fail(f"Unknown solver backend: {backend}. Available: {available}")

    override_solver = request.config.getini("override_solver")
    if override_solver:
        try:
            algo_name, impl_name = override_solver.split(":", 1)
        except ValueError:
            raise ValueError(
                f"override_solver must be in format 'algo:implementation', got: {override_solver}"
            )
        _solver_registry_to_use[algo_name] = getattr(nmo.solvers, impl_name)

    # save the original registry so that we can restore it after
    original = nmo.solvers.solver_registry.copy()
    nmo.solvers.solver_registry.clear()
    nmo.solvers.solver_registry.update(_solver_registry_to_use)

    try:
        yield
    finally:
        nmo.solvers.solver_registry.clear()
        nmo.solvers.solver_registry.update(original)


def pytest_addoption(parser):
    """Register custom ini options."""
    parser.addini("solver_backend", "Solver backend to use")
    parser.addini("override_solver", "Override solver as 'algorithm:implementation'")


@pytest.fixture
def gaussianGLM_model_instantiation():
    """Set up a Gaussian GLM for testing purposes.

    This fixture initializes a Gaussian GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (numpy.ndarray): Simulated input data.
            - np.random.normal(rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.PoissonGLM): Initialized model instance.
            - GLMParams(w_true, b_true): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of response.
    """
    np.random.seed(123)
    X = np.random.normal(size=(100, 5))
    b_true = np.zeros((1,))
    w_true = np.random.normal(size=(5,))
    observation_model = nmo.observation_models.GaussianObservations()
    regularizer = nmo.regularizer.UnRegularized()
    model = nmo.glm.GLM(
        observation_model, regularizer=regularizer, solver_name="LBFGS"
    )  # , solver_kwargs={"tol":1e-12})
    model.scale_ = 1.0
    rate = jax.numpy.einsum("k,tk->t", w_true, X) + b_true
    return X, np.random.normal(rate), model, GLMParams(w_true, b_true), rate


@pytest.fixture
def population_gaussianGLM_model_instantiation():
    """Set up a Population Gaussian GLM for testing purposes.

    This fixture initializes a Population Gaussian GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (numpy.ndarray): Simulated input data.
            - np.random.normal(rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.PoissonGLM): Initialized model instance.
            - GLMParams(w_true, b_true): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of response.
    """
    np.random.seed(123)
    X = np.random.normal(size=(200, 5)) * 10
    b_true = np.zeros((3,))
    w_true = np.random.normal(size=(5, 3))
    observation_model = nmo.observation_models.GaussianObservations()
    regularizer = nmo.regularizer.UnRegularized()
    model = nmo.glm.PopulationGLM(
        observation_model=observation_model,
        regularizer=regularizer,
        solver_name="LBFGS",
    )
    model.scale_ = 1.0
    rate = jax.numpy.einsum("ki,tk->ti", w_true, X) + b_true
    return X, np.random.normal(rate), model, GLMParams(w_true, b_true), rate


@pytest.fixture
def gaussianGLM_model_instantiation_pytree(gaussianGLM_model_instantiation):
    """Set up a Gaussian GLM for testing purposes.

    This fixture initializes a Gaussian GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (FeaturePytree): Simulated input data.
            - np.random.normal(rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.PoissonGLM): Initialized model instance.
            - GLMParams(w_true, b_true): True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of response.
    """
    X, spikes, model, true_params, rate = gaussianGLM_model_instantiation
    X_tree = nmo.pytrees.FeaturePytree(input_1=X[..., :3], input_2=X[..., 3:])
    true_params_tree = GLMParams(
        dict(input_1=true_params.coef[:3], input_2=true_params.coef[3:]),
        true_params.intercept,
    )
    model_tree = nmo.glm.GLM(
        model.observation_model, regularizer=model.regularizer, solver_name="LBFGS"
    )  # , solver_kwargs={"tol":1e-12})
    return X_tree, spikes, model_tree, true_params_tree, rate


@pytest.fixture
def population_gaussianGLM_model_instantiation_pytree(
    population_gaussianGLM_model_instantiation,
):
    """Set up a Population Gaussian GLM for testing purposes.

    This fixture initializes a Population Gaussian GLM with random parameters, simulates its response, and
    returns the test data, expected output, the model instance, true parameters, and the rate
    of response.

    Returns:
        tuple: A tuple containing:
            - X (FeaturePytree): Simulated input data.
            - np.random.normal(rate) (numpy.ndarray): Simulated spike responses.
            - model (nmo.glm.PoissonGLM): Initialized model instance.
            - GLMParams(w_true, b_true) : True weight and bias parameters.
            - rate (jax.numpy.ndarray): Simulated rate of response.
    """
    X, spikes, model, true_params, rate = population_gaussianGLM_model_instantiation
    X_tree = nmo.pytrees.FeaturePytree(input_1=X[..., :3], input_2=X[..., 3:])
    true_params_tree = GLMParams(
        dict(input_1=true_params.coef[:3], input_2=true_params.coef[3:]),
        true_params.intercept,
    )
    model_tree = nmo.glm.PopulationGLM(
        observation_model=model.observation_model,
        regularizer=model.regularizer,
        solver_name="LBFGS",
    )
    return X_tree, np.random.normal(rate), model_tree, true_params_tree, rate
