# mypy: disable-error-code=safe-super
from __future__ import annotations

import dataclasses as dc
from abc import ABC
from abc import abstractmethod
from typing import NoReturn

import numpy as np
import pandas as pd
import pyarrow as pa
from typeguard import check_type
from typeguard import TypeCheckError

from ..dtypes import Dtype
from ..output import check_output_format
from ..output import OutputFormat
from .typing import BatchData
from .typing import IsBatch
from .typing import NPArray
from .typing import NPMatrix
from .typing import NPSeqLens
from .typing import NPTensor
from .typing import PadMask
from .typing import SeqLens


class BatchMixinBase(ABC):
    @abstractmethod
    def as_format(
        self, output_format: OutputFormat, pin_memory: bool = False,
    ) -> IsBatch | NoReturn:
        raise ValueError(f'Unknown format: {output_format}')

    @abstractmethod
    def select(self, features: list[str]) -> IsBatch | NoReturn:
        raise ValueError('Batch data attribute is in an unknown format')

    @property
    @abstractmethod
    def pad_mask(self) -> PadMask | NoReturn:
        raise ValueError('Unknown format for "seq_lens" attribute')

    @abstractmethod
    def unpack(self, pad_value: int | float = np.nan) -> IsBatch | NoReturn:
        raise ValueError('Unknown format for "data" and "seq_lens" attributes')

    @property
    def is_packed(self: IsBatch) -> bool:
        if isinstance(self.data, (pa.Table, pa.Tensor)):
            return isinstance(self.data, pa.Table)
        else:
            return self.data.ndim == 2


class NumpyBatchMixin(BatchMixinBase):
    def as_format(
        self: IsBatch, output_format: OutputFormat, pin_memory: bool = False,
    ) -> IsBatch | NoReturn:
        if check_output_format(output_format, category='np'):
            return self
        else:
            return super().as_format(output_format=output_format, pin_memory=pin_memory)

    def select(self: IsBatch, features: list[str]) -> IsBatch | NoReturn:
        try:
            np_data = check_type(self.data, NPArray)
            feature_idx = self._features_to_idx(features)
            np_data = np_data[..., feature_idx]
            dtypes = [self.dtypes[i] for i in feature_idx]
            return dc.replace(self, data=np_data, dtypes=dtypes, feature_names=features)
        except TypeCheckError:
            return super().select(features=features)

    @property
    def pad_mask(self: IsBatch) -> PadMask | NoReturn:
        try:
            seq_lens = check_type(self.seq_lens, NPSeqLens)
            return np.arange(seq_lens.max()) >= seq_lens[:, None]
        except TypeCheckError:
            return super().pad_mask

    def unpack(
        self: IsBatch, pad_value: int | float = np.nan,
    ) -> IsBatch | NoReturn:
        def pad_packed(data: BatchData, seq_lens: SeqLens) -> NPTensor:
            np_data = check_type(data, NPMatrix)
            if np.issubdtype(np_data.dtype, np.integer):
                # Convert to float dtype so we can pad with np.nan
                np_data = np_data.astype('float')  # type: ignore[assignment]
            np_seq_lens = check_type(seq_lens, NPSeqLens)

            padded: list[NPMatrix] = []
            max_seq_len = int(np_seq_lens.max())
            for idx in self.iter_packed(seq_lens=np_seq_lens):
                x = np_data[idx, :]
                pad_len = max_seq_len - x.shape[0]
                if pad_len:
                    pad_width = ((0, pad_len), (0, 0))
                    x = np.pad(x, pad_width=pad_width, constant_values=pad_value)  # type: ignore[assignment]  # noqa: E501
                padded.append(x)
            return np.stack(padded, axis=0)

        try:
            tensor_data = pad_packed(data=self.data, seq_lens=self.seq_lens)
            tensor_targets: NPTensor | None = None
            if self.targets is not None:
                assert self.target_seq_lens is not None
                tensor_targets = pad_packed(
                    data=self.targets, seq_lens=self.target_seq_lens,
                )
            return dc.replace(self, data=tensor_data, targets=tensor_targets)
        except TypeCheckError:
            return super().unpack(pad_value=pad_value)


class PandasBatchMixin(BatchMixinBase):
    def as_format(
        self: IsBatch, output_format: OutputFormat, pin_memory: bool = False,
    ) -> IsBatch | NoReturn:
        if check_output_format(output_format, category='pd'):
            if self.is_packed:
                def set_df_column_dtypes(
                    df: pd.DataFrame, columns: list[str], dtypes: list[Dtype],
                ) -> pd.DataFrame:
                    for col, dtype in zip(columns, dtypes):
                        if (
                            df[col].isnull().any()
                            and isinstance(dtype, np.dtype)
                            and np.issubdtype(dtype, np.integer)
                        ):
                            # Replace NaN's with pd.NA and use the extension dtype Int64
                            df[col] = df[col].fillna(pd.NA).astype(dtype='Int64')
                        else:
                            df[col] = df[col].astype(dtype)
                    return df

                # Set the column types for the data DF
                data_df = pd.DataFrame(self.data, columns=self.feature_names)
                data_df = set_df_column_dtypes(
                    df=data_df, columns=self.feature_names, dtypes=self.dtypes,
                )
                # Set the columns types for the targets DF
                targets_df: pd.DataFrame | None = None
                if self.targets is not None:
                    assert self.target_names is not None and self.target_dtypes is not None  # noqa: E501
                    targets_df = pd.DataFrame(self.targets, columns=self.target_names)
                    targets_df = set_df_column_dtypes(
                        df=targets_df, columns=self.target_names,
                        dtypes=self.target_dtypes,
                    )
                return dc.replace(self, data=data_df, targets=targets_df)
            else:
                raise ValueError('Batch must be in packed form to convert to Pandas')
        else:
            return super().as_format(output_format=output_format, pin_memory=pin_memory)

    def select(self: IsBatch, features: list[str]) -> IsBatch | NoReturn:
        try:
            df = check_type(self.data, pd.DataFrame)
            feature_idx = self._features_to_idx(features)
            df = df.iloc[:, feature_idx]
            dtypes = [self.dtypes[i] for i in feature_idx]
            return dc.replace(self, data=df, dtypes=dtypes, feature_names=features)
        except TypeCheckError:
            return super().select(features=features)

    @property
    def pad_mask(self: IsBatch) -> PadMask | NoReturn:
        try:
            check_type(self.data, pd.DataFrame)
            raise ValueError('"pad_mask" is not applicable to packed data')
        except TypeCheckError:
            return super().pad_mask

    def unpack(
        self: IsBatch, pad_value: int | float = np.nan,
    ) -> IsBatch | NoReturn:
        try:
            check_type(self.data, pd.DataFrame)
            raise ValueError('Cannot convert Pandas Dataframe to tensor form...')
        except TypeCheckError:
            return super().unpack(pad_value=pad_value)


class PyArrowBatchMixin(BatchMixinBase):
    def as_format(
        self: IsBatch, output_format: OutputFormat, pin_memory: bool = False,
    ) -> IsBatch | NoReturn:
        if check_output_format(output_format, category='pa'):
            data: pa.Tensor | pa.Table
            if not self.is_packed and not self.is_numeric:
                raise ValueError('Batch must be numeric for Pyarrow Tensor format')

            if self.is_packed:
                # Build the Table for each batch row from arrays of feature columns
                pa_arrays = [
                    pa.array(self.data[:, d]) for d in range(self.feature_size)
                ]
                data = pa.Table.from_arrays(pa_arrays, names=self.feature_names)
            else:
                data = pa.Tensor.from_numpy(self.data)

            targets: pa.Tensor | pa.Table | None = None
            if self.targets is not None and self.target_names is not None:
                target_size = self.targets.shape[-1]
                if self.is_packed:
                    # Build the Table for each batch row from arrays of feature columns
                    pa_arrays = [
                        pa.array(self.targets[:, d]) for d in range(target_size)
                    ]
                    targets = pa.Table.from_arrays(pa_arrays, names=self.target_names)
                else:
                    targets = pa.Tensor.from_numpy(self.targets)

            return dc.replace(self, data=data, targets=targets)
        else:
            return super().as_format(output_format=output_format, pin_memory=pin_memory)

    def select(self: IsBatch, features: list[str]) -> IsBatch | NoReturn:
        try:
            tensor = check_type(self.data, pa.Tensor)
            np_data = tensor.to_numpy()
            feature_idx = self._features_to_idx(features)
            tensor = pa.Tensor.from_numpy(np_data[..., feature_idx])
            dtypes = [self.dtypes[i] for i in feature_idx]
            return dc.replace(self, data=tensor, dtypes=dtypes, feature_names=features)
        except TypeCheckError:
            ...
        try:
            table = check_type(self.data, pa.Table)
            table = table.select(features)
            feature_idx = self._features_to_idx(features)
            dtypes = [self.dtypes[i] for i in feature_idx]
            return dc.replace(self, data=table, dtypes=dtypes, feature_names=features)
        except TypeCheckError:
            ...
        return super().select(features=features)

    @property
    def pad_mask(self: IsBatch) -> PadMask | NoReturn:
        try:
            seq_lens = check_type(self.seq_lens, NPSeqLens)
            return np.arange(seq_lens.max()) >= seq_lens[:, None]
        except TypeCheckError:
            return super().pad_mask

    def unpack(
        self: IsBatch, pad_value: int | float = np.nan,
    ) -> IsBatch | NoReturn:
        def pad_packed(data: BatchData, seq_lens: SeqLens) -> pa.Tensor:
            pa_data = check_type(data, pa.Table)
            np_seq_lens = check_type(seq_lens, NPSeqLens)
            if not self.is_numeric:
                raise ValueError('Batch must be numeric for Pyarrow Tensor format')

            padded: list[NPMatrix] = []
            np_data = pa_data.to_pandas().to_numpy().astype(float)
            max_seq_len = np_seq_lens.max()
            for idx in self.iter_packed(seq_lens=np_seq_lens):
                x = np_data[idx, :]
                pad_len = max_seq_len - x.shape[0]
                if pad_len:
                    pad_width = ((0, pad_len), (0, 0))
                    x = np.pad(x, pad_width=pad_width, constant_values=pad_value)
                padded.append(x)
            return pa.Tensor.from_numpy(np.stack(padded, axis=0))
        try:
            tensor_data = pad_packed(data=self.data, seq_lens=self.seq_lens)
            tensor_targets: pa.Tensor | None = None
            if self.targets is not None:
                assert self.target_seq_lens is not None
                tensor_targets = pad_packed(
                    data=self.targets, seq_lens=self.target_seq_lens,
                )
            return dc.replace(self, data=tensor_data, targets=tensor_targets)
        except TypeCheckError:
            return super().unpack(pad_value=pad_value)


class TorchBatchMixin(BatchMixinBase):
    def as_format(
        self: IsBatch, output_format: OutputFormat, pin_memory: bool = False,
    ) -> IsBatch | NoReturn:
        if check_output_format(output_format, category='pt'):
            from tempora.utils.batch.typing.torch import torch
            from tempora.utils.batch.typing.torch import TorchArray
            from tempora.utils.batch.typing.torch import TorchSeqLens

            # We want to ensure a memcopy here so can't use torch.from_numpy()
            data = torch.tensor(self.data, pin_memory=pin_memory)
            seq_lens = torch.tensor(
                self.seq_lens.astype('int64'), dtype=torch.int64, pin_memory=pin_memory,
            )
            targets: TorchArray | None = None
            target_seq_lens: TorchSeqLens | None = None
            if self.targets is not None:
                assert self.target_seq_lens is not None
                targets = torch.tensor(self.targets, pin_memory=pin_memory)
                target_seq_lens = torch.tensor(
                    self.target_seq_lens.astype('int64'), dtype=torch.int64,
                    pin_memory=pin_memory,
                )
            return dc.replace(
                self, data=data, seq_lens=seq_lens,
                targets=targets, target_seq_lens=target_seq_lens,
            )
        else:
            return super().as_format(output_format=output_format, pin_memory=pin_memory)

    def select(self: IsBatch, features: list[str]) -> IsBatch | NoReturn:
        try:
            from tempora.utils.batch.typing.torch import TorchArray

            torch_data = check_type(self.data, TorchArray)
            feature_idx = self._features_to_idx(features)
            torch_data = torch_data[..., feature_idx]
            dtypes = [self.dtypes[i] for i in feature_idx]
            return dc.replace(
                self, data=torch_data, dtypes=dtypes, feature_names=features,
            )
        except (ImportError, TypeCheckError):
            return super().select(features=features)

    @property
    def pad_mask(self: IsBatch) -> PadMask | NoReturn:
        try:
            from tempora.utils.batch.typing.torch import torch
            from tempora.utils.batch.typing.torch import TorchSeqLens

            seq_lens = check_type(self.seq_lens, TorchSeqLens)
            max_seq_len = int(seq_lens.max())
            return torch.arange(max_seq_len) >= seq_lens[:, None]
        except (ImportError, TypeCheckError):
            return super().pad_mask

    def unpack(
        self: IsBatch, pad_value: int | float = np.nan,
    ) -> IsBatch | NoReturn:
        def pad_packed(data: BatchData, seq_lens: SeqLens) -> TorchTensor:
            torch_data = check_type(data, TorchMatrix)
            torch_seq_lens = check_type(seq_lens, TorchSeqLens)
            from torch.nn.functional import pad
            padded: list[TorchMatrix] = []
            max_seq_len = int(torch_seq_lens.max())
            for idx in self.iter_packed(seq_lens=torch_seq_lens):
                x = torch_data[idx, :]
                pad_len = max_seq_len - x.shape[0]
                if pad_len:
                    x = pad(x, pad=[0, 0, 0, pad_len], value=pad_value)
                padded.append(x)
            return torch.stack(padded, dim=0)

        try:
            from tempora.utils.batch.typing.torch import torch
            from tempora.utils.batch.typing.torch import TorchMatrix
            from tempora.utils.batch.typing.torch import TorchSeqLens
            from tempora.utils.batch.typing.torch import TorchTensor

            if np.isnan(pad_value):
                pad_value = torch.nan

            tensor_data = pad_packed(data=self.data, seq_lens=self.seq_lens)
            tensor_targets: TorchTensor | None = None
            if self.targets is not None:
                assert self.target_seq_lens is not None
                tensor_targets = pad_packed(
                    data=self.targets, seq_lens=self.target_seq_lens,
                )
            return dc.replace(self, data=tensor_data, targets=tensor_targets)
        except TypeCheckError:
            return super().unpack(pad_value=pad_value)


class TFBatchMixin(BatchMixinBase):
    def as_format(
        self: IsBatch, output_format: OutputFormat, pin_memory: bool = False,
    ) -> IsBatch | NoReturn:
        if check_output_format(output_format, category='tf'):
            from tempora.utils.batch.typing.tensorflow import tf
            from tempora.utils.batch.typing.tensorflow import TFArray
            from tempora.utils.batch.typing.tensorflow import TFSeqLens

            data = tf.convert_to_tensor(self.data)
            seq_lens = tf.convert_to_tensor(self.seq_lens)
            targets: TFArray | None = None
            target_seq_lens: TFSeqLens | None = None
            if self.targets is not None:
                assert self.target_seq_lens is not None
                targets = tf.convert_to_tensor(self.targets)
                target_seq_lens = tf.convert_to_tensor(self.target_seq_lens)
            return dc.replace(
                self, data=data, seq_lens=seq_lens,
                targets=targets, target_seq_lens=target_seq_lens,
            )
        else:
            return super().as_format(output_format=output_format, pin_memory=pin_memory)

    def select(self: IsBatch, features: list[str]) -> IsBatch | NoReturn:
        try:
            from tempora.utils.batch.typing.tensorflow import tf
            from tempora.utils.batch.typing.tensorflow import TFArray

            tf_data = check_type(self.data, TFArray)
            feature_idx = self._features_to_idx(features)
            axis = 1 if self.is_packed else 2
            tf_data = tf.gather(tf_data, indices=feature_idx, axis=axis)
            dtypes = [self.dtypes[i] for i in feature_idx]
            return dc.replace(
                self, data=tf_data, dtypes=dtypes, feature_names=features,
            )
        except (ImportError, TypeCheckError):
            return super().select(features=features)

    @property
    def pad_mask(self: IsBatch) -> PadMask | NoReturn:
        try:
            from tempora.utils.batch.typing.tensorflow import tf
            from tempora.utils.batch.typing.tensorflow import TFSeqLens

            seq_lens = check_type(self.seq_lens, TFSeqLens)
            max_seq_len = tf.math.reduce_max(seq_lens)
            from tf.experimental.numpy import arange as tf_arange  # noqa
            return tf_arange(max_seq_len) >= seq_lens[:, None]
        except (ImportError, TypeCheckError):
            return super().pad_mask

    def unpack(
        self: IsBatch, pad_value: int | float = np.nan,
    ) -> IsBatch | NoReturn:
        def pad_packed(data: BatchData, seq_lens: SeqLens) -> TFTensor:
            tf_data = check_type(data, TFMatrix)
            tf_seq_lens = check_type(seq_lens, TFSeqLens)
            padded: list[TFMatrix] = []
            max_seq_len = int(tf.math.reduce_max(tf_seq_lens))
            for idx in self.iter_packed(seq_lens=tf_seq_lens):
                x = tf_data[idx, :]
                pad_len = max_seq_len - x.shape[0]
                if pad_len:
                    paddings = tf.constant([[0, pad_len], [0, 0]])
                    x = tf.pad(x, paddings=paddings, constant_values=pad_value)
                padded.append(x)
            return tf.stack(padded, axis=0)

        try:
            from tempora.utils.batch.typing.tensorflow import tf
            from tempora.utils.batch.typing.tensorflow import TFMatrix
            from tempora.utils.batch.typing.tensorflow import TFSeqLens
            from tempora.utils.batch.typing.tensorflow import TFTensor

            tensor_data = pad_packed(data=self.data, seq_lens=self.seq_lens)
            tensor_targets: TFTensor | None = None
            if self.targets is not None:
                assert self.target_seq_lens is not None
                tensor_targets = pad_packed(
                    data=self.targets, seq_lens=self.target_seq_lens,
                )
            return dc.replace(self, data=tensor_data, targets=tensor_targets)
        except TypeCheckError:
            return super().unpack(pad_value=pad_value)


class JaxBatchMixin(BatchMixinBase):
    def as_format(
        self: IsBatch, output_format: OutputFormat, pin_memory: bool = False,
    ) -> IsBatch | NoReturn:
        if check_output_format(output_format, category='jx'):
            from tempora.utils.batch.typing.jax import jnp
            from tempora.utils.batch.typing.jax import JaxArray
            from tempora.utils.batch.typing.jax import JaxSeqLens

            data = jnp.asarray(self.data)
            seq_lens = jnp.asarray(self.seq_lens)
            targets: JaxArray | None = None
            target_seq_lens: JaxSeqLens | None = None
            if self.targets is not None:
                assert self.target_seq_lens is not None
                targets = jnp.asarray(self.targets)
                target_seq_lens = jnp.asarray(self.target_seq_lens)
            return dc.replace(
                self, data=data, seq_lens=seq_lens,
                targets=targets, target_seq_lens=target_seq_lens,
            )
        else:
            return super().as_format(output_format=output_format, pin_memory=pin_memory)

    def select(
        self: IsBatch, features: list[str],
    ) -> IsBatch | NoReturn:
        try:
            from tempora.utils.batch.typing.jax import JaxArray

            jax_data = check_type(self.data, JaxArray)
            feature_idx = self._features_to_idx(features)
            jax_data = jax_data[..., feature_idx]
            dtypes = [self.dtypes[i] for i in feature_idx]
            return dc.replace(
                self, data=jax_data, dtypes=dtypes, feature_names=features,
            )
        except (ImportError, TypeCheckError):
            return super().select(features=features)

    @property
    def pad_mask(self: IsBatch) -> PadMask | NoReturn:
        try:
            from tempora.utils.batch.typing.jax import jnp
            from tempora.utils.batch.typing.jax import JaxSeqLens

            seq_lens = check_type(self.seq_lens, JaxSeqLens)
            max_seq_len = jnp.max(seq_lens)
            return jnp.arange(max_seq_len) >= seq_lens[:, None]
        except (ImportError, TypeCheckError):
            return super().pad_mask

    def unpack(
        self: IsBatch, pad_value: int | float = np.nan,
    ) -> IsBatch | NoReturn:
        def pad_packed(data: BatchData, seq_lens: SeqLens) -> JaxTensor:
            jax_data = check_type(data, JaxMatrix)
            if np.issubdtype(jax_data.dtype, np.integer):
                # Convert to float dtype so we can pad with np.nan
                jax_data = jax_data.astype('float')
            jax_seq_lens = check_type(seq_lens, JaxSeqLens)

            padded: list[JaxMatrix] = []
            max_seq_len = jnp.max(jax_seq_lens)
            for idx in self.iter_packed(seq_lens=jax_seq_lens):
                x = jax_data[idx, :]
                pad_len = max_seq_len - x.shape[0]
                if pad_len:
                    pad_width = ((0, pad_len), (0, 0))
                    x = jnp.pad(x, pad_width=pad_width, constant_values=pad_value)
                padded.append(x)
            return jnp.stack(padded, axis=0)

        try:
            from tempora.utils.batch.typing.jax import jnp
            from tempora.utils.batch.typing.jax import JaxMatrix
            from tempora.utils.batch.typing.jax import JaxSeqLens
            from tempora.utils.batch.typing.jax import JaxTensor

            tensor_data = pad_packed(data=self.data, seq_lens=self.seq_lens)
            tensor_targets: JaxTensor | None = None
            if self.targets is not None:
                assert self.target_seq_lens is not None
                tensor_targets = pad_packed(
                    data=self.targets, seq_lens=self.target_seq_lens,
                )
            return dc.replace(self, data=tensor_data, targets=tensor_targets)
        except TypeCheckError:
            return super().unpack(pad_value=pad_value)
