"""Collection of nemos typing."""

from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    NamedTuple,
    Tuple,
    TypeAlias,
    TypeVar,
)

import jax.numpy as jnp
from jax.typing import ArrayLike
from numpy.typing import NDArray

if TYPE_CHECKING:
    import pynapple as nap

Pytree: TypeAlias = Any
Params: TypeAlias = Pytree
Aux = TypeVar("Aux")
SolverState = TypeVar("SolverState")
StepResult: TypeAlias = Tuple[Params, SolverState, Aux]
DESIGN_INPUT_TYPE = "Union[jnp.ndarray, FeaturePytree, nap.TsdFrame]"

# copying jax.random's annotation
KeyArrayLike = ArrayLike

SolverRun = Callable[
    [
        Params,  # parameters, could be any pytree
        jnp.ndarray,  # Predictors (i.e. model design for GLM)
        jnp.ndarray,
    ],  # Output (neural activity)
    StepResult,
]

SolverInit = Callable[
    [
        Params,  # parameters, could be any pytree
        jnp.ndarray,  # Predictors (i.e. model design for GLM)
        jnp.ndarray,
    ],  # Output (neural activity)
    SolverState,
]

SolverUpdate = Callable[
    [
        Params,  # parameters, could be any pytree
        NamedTuple,
        jnp.ndarray,  # Predictors (i.e. model design for GLM)
        jnp.ndarray,
    ],  # Output (neural activity)
    StepResult,
]

ProximalOperator = Callable[
    [
        Params,  # parameters, could be any pytree
        float,  # Regularizer strength (for now float, eventually pytree)
        float,
    ],  # Step-size for optimization (must be a float)
    Tuple[jnp.ndarray, jnp.ndarray],
]

FeatureMatrix: TypeAlias = "nap.TsdFrame | NDArray"

# User provided init_params (e.g. for GLMs Tuple[array, array])
UserProvidedParamsT = TypeVar("UserProvidedParamsT")
# Model internal representation (e.g. for GLMs nemos.glm.glm.GLMParams)
ModelParamsT = TypeVar("ModelParamsT")
