from __future__ import annotations

import dataclasses as dc
from typing import Iterator
from typing import NoReturn
from typing import Protocol
from typing import TYPE_CHECKING
from typing import TypeVar
from typing import Union

import jaxtyping as jxt
import numpy as np
import pandas as pd
import pyarrow as pa

from ...dtypes import Dtype
from ...output import OutputFormat
from ..metadata import BatchMetadata


# Numpy types
NPArray = jxt.Shaped[np.ndarray, '...']  # np array dtype can be 'O' if we have strings
NPMatrix = jxt.Shaped[np.ndarray, 'Time Features']
NPTensor = jxt.Shaped[np.ndarray, 'Batch Time Features']
NPPadMask = jxt.Bool[np.ndarray, 'Batch Time']
NPSeqLens = jxt.UInt[np.ndarray, 'Batch']


# ------------------------------------------------------------------------------------
# Type Alias def's for static type checks (and runtime imports) vs runtime type checks
# ------------------------------------------------------------------------------------
# For PyTorch / Tensorflow / Jax we only try to import the actual package, to define the
# relevant type aliases, if we are type checking. At run-time we always import from the
# mock types (np array's) so that we can avoid the slow import times that can occur when
# importing tensorflow or pytorch for the first time. This helps to speed up imports
# such as - "from tempora.samplers import *"
#
# For runtime type *checks* with typeguard's check_type function, e.g. for various mixin
# operations within the BatchData class, we need the proper type aliases derived from
# the relevant pytorch/tf/jax data types, and not the mock type aliases. These are
# defined within the jax/torch/tensorflow.py modules under the batch_typing directory.
#
# NB: we cannot combine the type alias logic for static type checking / runtime imports
# with the module-level type alias def's for runtime type checks, since mypy does not
# support "conditional type aliases", see:
# https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases

# Pytorch types
if TYPE_CHECKING:
    try:
        import torch
    except ImportError:
        from . import mock_types as torch  # type: ignore[no-redef]
else:
    from . import mock_types as torch

TorchArray = jxt.Float[torch.Tensor, '...']
TorchMatrix = jxt.Float[torch.Tensor, 'Time Features']
TorchTensor = jxt.Float[torch.Tensor, 'Batch Time Features']
TorchPadMask = jxt.Bool[torch.Tensor, 'Batch Time']
TorchSeqLens = jxt.Int[torch.Tensor, 'Batch']  # Torch only supports 8-bit unsigned ints

# Tensorflow types
if TYPE_CHECKING:
    try:
        import tensorflow as tf
    except ImportError:
        from . import mock_types as tf
else:
    from . import mock_types as tf

TFArray = jxt.Float[tf.Tensor, '...']
TFMatrix = jxt.Float[tf.Tensor, 'Time Features']
TFTensor = jxt.Float[tf.Tensor, 'Batch Time Features']
TFPadMask = jxt.Bool[tf.Tensor, 'Batch Time']
TFSeqLens = jxt.UInt[tf.Tensor, 'Batch']

# Jax types
if TYPE_CHECKING:
    try:
        import jax
    except ImportError:
        # 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:
    from . import mock_types as jax
    jxt.Array = jax.Array

JaxArray = jxt.Float[jxt.Array, '...']
JaxMatrix = jxt.Float[jxt.Array, 'Time Features']
JaxTensor = jxt.Float[jxt.Array, 'Batch Time Features']
JaxPadMask = jxt.Bool[jxt.Array, 'Batch Time']
JaxSeqLens = jxt.UInt[jxt.Array, 'Batch']


BatchDataPacked = TypeVar(
    'BatchDataPacked', NPMatrix, TorchMatrix, TFMatrix, JaxMatrix, pa.Table,
    pd.DataFrame,
)
BatchDataTensor = TypeVar(
    'BatchDataTensor', NPTensor, TorchTensor, TFTensor, JaxTensor, pa.Tensor,
)
BatchData = TypeVar(
    'BatchData', NPTensor, TorchTensor, TFTensor, JaxTensor, pa.Tensor,
    NPMatrix, TorchMatrix, TFMatrix, JaxMatrix, pa.Table, pd.DataFrame,
)
SeqLens = TypeVar('SeqLens', NPSeqLens, TorchSeqLens, TFSeqLens, JaxSeqLens)

# Mutable batch data types for which we can pre-allocate memory for feature computation.
# We only consider types that are both mutable and arrays: we exclude Pandas Dataframes
# to avoid having to deal with .loc indexing as a special case throughout the code base.
# NB: Jax arrays, TensorFlow tensors, PyArrow Tables and Tensors are all immutable.
MutableBatchData = TypeVar(
    'MutableBatchData', NPTensor, TorchTensor, NPMatrix, TorchMatrix,
)
MutableSeqLens = TypeVar('MutableSeqLens', NPSeqLens, TorchSeqLens)

PadMask = Union[NPPadMask, TorchPadMask, TFPadMask, JaxPadMask]


@dc.dataclass(frozen=True)
class IsBatch(Protocol[BatchData, SeqLens]):
    data: BatchData
    seq_lens: SeqLens
    feature_names: list[str]
    dtypes: list[Dtype]
    targets: BatchData | None = None
    target_seq_lens: SeqLens | None = None
    target_names: list[str] | None = None
    target_dtypes: list[Dtype] | None = None
    metadata: BatchMetadata | None = None

    def as_format(
        self, output_format: OutputFormat, pin_memory: bool,
    ) -> IsBatch[BatchData, SeqLens] | NoReturn: ...

    def select(
        self, features: list[str],
    ) -> IsBatch[BatchData, SeqLens] | NoReturn: ...

    def _get_feature_idx(self, feature: str) -> int | NoReturn: ...

    def _features_to_idx(self, features: list[str]) -> list[int]: ...

    def unpack(
        self, pad_value: int | float = np.nan,
    ) -> IsBatch[BatchData, SeqLens] | NoReturn: ...

    def to_tensor(
        self, pad_value: int | float = np.nan,
    ) -> IsBatch[BatchData, SeqLens] | NoReturn: ...

    def iter_packed(
        self, seq_lens: SeqLens | None = None,
    ) -> Iterator[slice] | NoReturn: ...

    @property
    def is_packed(self) -> bool: ...  # noqa

    @property
    def pad_mask(self) -> PadMask | NoReturn: ...  # noqa

    @property
    def columns(self) -> list[str] | None: ...  # noqa

    @property
    def feature_size(self) -> int: ...  # noqa

    @property
    def is_numeric(self) -> bool: ...  # noqa
