from __future__ import annotations

from typing import TYPE_CHECKING

import jaxtyping as jxt
try:
    import jax
    import jax.numpy as jnp  # noqa: F401
except ImportError:
    if TYPE_CHECKING:
        # Create the "Array" attribute in the jaxtyping module. This is usually an alias
        # for jax.numpy.ndarray but would not be present if jax is not installed,
        # resulting in an AttributeError in the Jax* definitions below
        from . import mock_types as jax
        jxt.Array = jax.Array
    else:
        raise ImportError(
            'jax not installed - try one of the following:\n'
            '-> with pip: pip install tempora-ai[jax]\n'
            '-> with poetry: poetry add tempora-ai -E jax',
        )


# Jax type aliases for runtime type *checks* (e.g. with typeguard's check_type function)
JaxArray = jxt.Float[jax.Array, '...']
JaxMatrix = jxt.Float[jax.Array, 'Time Features']
JaxTensor = jxt.Float[jax.Array, 'Batch Time Features']
JaxPadMask = jxt.Bool[jax.Array, 'Batch Time']
JaxSeqLens = jxt.UInt[jax.Array, 'Batch']
