from __future__ import annotations

import dataclasses as dc
import json
import logging
from functools import partial
from itertools import islice
from pathlib import Path
from typing import Any
from typing import cast
from typing import Generic
from typing import Iterator
from typing import NoReturn

import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.flight as flight
import pyarrow.fs as fs
import pyarrow.parquet as pq
from typeguard import check_type
from typing_extensions import TypeGuard

from ..dtypes import Dtype
from ..fs import create_pyarrow_filesystem
from ..fs import remove_uri_prefix
from ..misc import to_list
from ..output import check_output_format
from ..output import get_format_name
from ..output import is_numeric_format
from ..output import OutputFormat
from .formats import JaxBatchMixin
from .formats import NumpyBatchMixin
from .formats import PandasBatchMixin
from .formats import PyArrowBatchMixin
from .formats import TFBatchMixin
from .formats import TorchBatchMixin
from .metadata import BatchMetadata
from .metadata import MetadataJSONEncoder
from .metadata import SegmentMetadata
from .typing import BatchData
from .typing import JaxMatrix
from .typing import JaxSeqLens
from .typing import JaxTensor
from .typing import MutableBatchData
from .typing import MutableSeqLens
from .typing import NPArray
from .typing import NPMatrix
from .typing import NPSeqLens
from .typing import NPTensor
from .typing import PadMask
from .typing import SeqLens
from .typing import TFMatrix
from .typing import TFSeqLens
from .typing import TFTensor
from .typing import TorchMatrix
from .typing import TorchSeqLens
from .typing import TorchTensor


log = logging.getLogger(__name__)


@dc.dataclass(frozen=True)
class TSSegment:
    data: pa.RecordBatch
    metadata: SegmentMetadata
    targets: pa.RecordBatch | None = None

    @classmethod
    def from_chunk(cls, chunk: flight.FlightStreamChunk) -> TSSegment | None:
        segment: TSSegment | None = None

        flight_data: pa.RecordBatch
        flight_metadata: pa.Buffer | None
        flight_data, flight_metadata = chunk

        if flight_data.num_rows:
            data: pa.RecordBatch
            targets: pa.RecordBatch | None = None

            if flight_data.column_names == ['__tempora_union']:
                # Data + targets: flight data schema is a union of two struct array
                # types, one defining the data schema and the other the targets schema.
                # The data rows have type code 0 and the target rows 1 in the UnionArray
                union_array: pa.UnionArray = flight_data.column(0)

                # Extract the data rows and target rows into individual RecordBatch's
                data = pa.RecordBatch.from_struct_array(union_array.field(0))
                targets = pa.RecordBatch.from_struct_array(union_array.field(1))
            else:
                # No targets - flight data recbatch *is* the sampled data recbatch
                data = flight_data

            if flight_metadata:
                segment_metadata = SegmentMetadata.from_flight_metadata(
                    metadata=flight_metadata,
                )
            else:
                segment_metadata = SegmentMetadata()

            segment = TSSegment(
                data=data, metadata=segment_metadata, targets=targets,
            )

        return segment


@dc.dataclass(frozen=True)
class Batch(
    NumpyBatchMixin, PandasBatchMixin, PyArrowBatchMixin, TorchBatchMixin,
    JaxBatchMixin, TFBatchMixin, Generic[BatchData, SeqLens],
):
    """Class for batches generated by a Sampler.

    Args:
        data: Batch data either in tensor or (packed) matrix form
        seq_lens: Length of each data sequence in the batch
        feature_names: Names of the batch features / columns
        dtypes: Batch data dtypes
        targets: Optional batch targets
        target_seq_lens: Length of each target sequence
        target_names: Names of the target features / columns
        target_dtypes: target dtypes
        metadata: Batch metadata
    """
    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 __post_init__(self) -> None:
        #######################################################################
        # Check that the shape of the various Batch attributes are compatible
        #######################################################################
        B = len(self.seq_lens)
        num_data_rows = int(self.data.shape[0])

        if self.is_packed:
            try:
                seq_lens_sum = self.seq_lens.sum()
            except AttributeError:
                # Handle Tensorflow tensors
                from tempora.utils.batch.typing.tensorflow import TFSeqLens  # noqa: E501
                from tensorflow.math import reduce_sum  # noqa
                seq_lens = check_type(self.seq_lens, TFSeqLens)
                seq_lens_sum = reduce_sum(seq_lens)

            if seq_lens_sum != num_data_rows:
                raise ValueError(
                    f'Sum of the seq_lens values ({seq_lens_sum}) should equal the '
                    f'number of rows of the packed data array ({num_data_rows})',
                )
        else:
            if B != num_data_rows:
                raise ValueError(
                    f'Size of seq_lens ({B}) should equal the data batch size '
                    f'{num_data_rows}',
                )

        K = len(self.feature_names)
        D = self.feature_size
        if K != D:
            raise ValueError(
                f'Number of feature names ({K}) should equal the feature size ({D})',
            )

        if self.targets is not None:
            assert self.target_seq_lens is not None
            BT = len(self.target_seq_lens)
            num_target_rows = int(self.targets.shape[0])

            if BT != B:
                raise ValueError(
                    f'Size of target_seq_lens ({BT}) should equal the batch size ({B})',
                )

            if self.is_packed:
                try:
                    seq_lens_sum = self.target_seq_lens.sum()
                except AttributeError:
                    # Handle Tensorflow tensors
                    from tempora.utils.batch.typing.tensorflow import TFSeqLens  # noqa: E501
                    from tensorflow.math import reduce_sum  # noqa
                    target_seq_lens = check_type(self.target_seq_lens, TFSeqLens)
                    seq_lens_sum = reduce_sum(target_seq_lens)

                if seq_lens_sum != num_target_rows:
                    raise ValueError(
                        f'Sum of the target_seq_lens values ({seq_lens_sum}) should '
                        f'equal the number of rows of the packed target data array '
                        f'({num_target_rows})',
                    )
            else:
                if BT != num_target_rows:
                    raise ValueError(
                        f'Size of target_seq_lens ({BT}) should equal the targets '
                        f'batch size {num_target_rows}',
                    )

            assert self.target_names is not None
            H = len(self.target_names)
            DT = int(self.targets.shape[-1])
            if H != DT:
                raise ValueError(
                    f'Number of target names ({H}) should equal the target size ({DT})',
                )

    @classmethod
    def _from_segments(
        cls, segments: list[TSSegment],
        as_tensor: bool = False,
        pad_value: int | float = np.nan,
        output_format: OutputFormat | None = None,
        pin_memory: bool = False,
    ) -> Batch[BatchData, SeqLens]:
        # Extract the data, targets and metadata from the segments list
        batch_data = [segment.data for segment in segments]
        feature_names = batch_data[0].schema.names
        batch_targets = [segment.targets for segment in segments]
        batch_metadata = [segment.metadata for segment in segments]

        # For numeric and string columns pyarrow's to_pandas_dtype() returns an object
        # of type "type" (e.g. "numpy.dtype[int64]"), however for timestamps it returns
        # an actual Pandas dtype ("pandas.core.dtypes.dtypes.DatetimeTZDtype"). This
        # helper function converts the raw types to actual numpy dtypes for everything
        # but timestamps. This also ensures a nicer __repr__output for the dtypes.
        def make_dtypes(schema: pa.Schema) -> list[Dtype]:
            dtypes: list[Dtype] = []
            for pa_dtype in schema.types:
                pd_dtype = pa_dtype.to_pandas_dtype()
                try:
                    dtypes.append(np.dtype(pd_dtype))
                except TypeError:
                    dtypes.append(pd_dtype)
            return dtypes

        # Create the Batch in *packed* numpy format
        np_data = np.concatenate([rec.to_pandas().to_numpy() for rec in batch_data])
        seq_lens = np.array([rec.num_rows for rec in batch_data], dtype=np.uint32)
        data_dtypes = make_dtypes(batch_data[0].schema)

        np_targets: np.ndarray | None = None
        target_names: list[str] | None = None
        target_seq_lens: np.ndarray | None = None
        target_dtypes: list[Dtype] | None = None

        def all_valid_targets(
            recs: list[pa.RecordBatch | None],
        ) -> TypeGuard[list[pa.RecordBatch]]:
            return all(isinstance(rec, pa.RecordBatch) for rec in recs)

        # Proceed if all segments have valid target data
        if all_valid_targets(batch_targets):
            # Target tables for all segments must have the same number of columns
            assert len({rec.num_columns for rec in batch_targets}) == 1

            # Unify the list of target record batches into a single np array
            np_targets = np.concatenate(
                [rec.to_pandas().to_numpy() for rec in batch_targets],
            )
            target_seq_lens = np.array(
                [rec.num_rows for rec in batch_targets], dtype=np.uint32,
            )
            target_names = batch_targets[0].schema.names
            target_dtypes = make_dtypes(batch_targets[0].schema)
        else:
            # If the batch targets are not all RecordBatches they should be all None's
            assert all(rec is None for rec in batch_targets)

        ts_batch = Batch[NPArray, NPSeqLens](
            data=np_data,
            seq_lens=seq_lens,
            feature_names=feature_names,
            dtypes=data_dtypes,
            targets=np_targets,
            target_seq_lens=target_seq_lens,
            target_names=target_names,
            target_dtypes=target_dtypes,
            metadata=batch_metadata,
        )

        if as_tensor:
            ts_batch = ts_batch.unpack(pad_value=pad_value)

        if output_format and not check_output_format(output_format, category='np'):
            return ts_batch.as_format(
                output_format=output_format, pin_memory=pin_memory,
            )
        else:
            return ts_batch

    @classmethod
    def read_parquet(
        cls, filepath: str | Path,
        filesystem: fs.FileSystem | None = None,
        fs_config: dict[str, Any] | None = None,
        as_tensor: bool = False,
        pad_value: int | float = np.nan,
        output_format: OutputFormat | None = None,
        pin_memory: bool = False,
    ) -> Batch[BatchData, SeqLens]:
        """Alias of from_parquet (for compatibility with the Pandas Dataframe API)."""
        return cls.from_parquet(
            filepath=filepath,
            filesystem=filesystem,
            fs_config=fs_config,
            as_tensor=as_tensor,
            pad_value=pad_value,
            output_format=output_format,
            pin_memory=pin_memory,
        )

    @classmethod
    def from_parquet(
        cls, filepath: str | Path,
        filesystem: fs.FileSystem | None = None,
        fs_config: dict[str, Any] | None = None,
        as_tensor: bool = False,
        pad_value: int | float = np.nan,
        output_format: OutputFormat | None = None,
        pin_memory: bool = False,
    ) -> Batch[BatchData, SeqLens]:
        """Deserialize a batch object from a parquet file.

        Args:
            filepath: Path to the serialized batch parquet file.
            filesystem: Optional [PyArrow filesystem](https://arrow.apache.org/docs/python/api/filesystems.html#filesystem-implementations) to read the parquet file from.
            fs_config: Custom arguments to pass to [PyArrow filesystem](https://arrow.apache.org/docs/python/api/filesystems.html#filesystem-implementations) if `filesystem` is not provided.
            as_tensor: Return the batch data in tensor form, otherwise (packed) matrix.
            pad_value: Tensor padding value to use (for uneven sequence lengths).
            output_format: Format to return the batch in.
            pin_memory: For PyTorch tensors use page-locked CPU memory to speed up data
                transfers to GPUs.

        Returns:
            Deserialized batch object.
        """  # noqa: E501
        # Read the parquet file into a Pyarrow Table then extract the schema metadata
        filepath = str(filepath)
        src_filepath = filepath  # make a copy of the filepath before rm'ing uri prefix
        if filesystem is None:
            filesystem = create_pyarrow_filesystem(source=filepath, fs_config=fs_config)
            filepath = remove_uri_prefix(filepath)

        # Reading parquet requires random access, so we need to read the stream into a
        # buffer before passing to pq.read_table. NB: this approach allows us to handle
        # gzipped files, unlike filesystem.open_input_file() which only works for .pq
        with filesystem.open_input_stream(filepath, compression='detect') as stream:
            buffer = pa.BufferReader(stream.read())
            table = pq.read_table(buffer)

        schema_metadata = table.schema.metadata
        try:
            seq_lens = json.loads(schema_metadata[b'__tempora_seq_lens'])
            feature_names = json.loads(schema_metadata[b'__tempora_feature_names'])
            target_seq_lens = json.loads(schema_metadata[b'__tempora_target_seq_lens'])
            target_names = json.loads(schema_metadata[b'__tempora_target_names'])
            batch_metadata = json.loads(schema_metadata[b'__tempora_batch_metadata'])
            batch_metadata = [SegmentMetadata(dct) for dct in batch_metadata]
            # Add the source filepath to the batch metadata if it's not present
            for segment_metadata in batch_metadata:
                if 'source' not in segment_metadata:
                    segment_metadata['source'] = src_filepath
        except KeyError as e:
            raise KeyError('Invalid parquet file - missing tempora metadata') from e

        df = table.to_pandas()
        # NB: must use [] column indexing rather than . to prevent python name mangling
        is_data_row = df['__tempora_batch_data']

        # Select the batch data rows and extract the dtypes from the field metadata
        data_df = df.loc[is_data_row, feature_names]
        np_seq_lens = np.array(seq_lens, dtype=np.uint32)
        # TODO: add support for PandasDtype in utils.dtype
        dtypes: list[Dtype] = [
            np.dtype(table.schema.field(name).metadata[b'dtype'])
            for name in feature_names
        ]

        if is_data_row.all():
            ts_batch = Batch[NPArray, NPSeqLens](
                data=data_df.to_numpy(),
                seq_lens=np_seq_lens,
                feature_names=feature_names,
                dtypes=dtypes,
                metadata=batch_metadata,
            )
        else:
            # Select the target data rows and extract the target dtypes
            targets_df = df.loc[~is_data_row, target_names]
            np_target_seq_lens = np.array(target_seq_lens, dtype=np.uint32)
            # TODO: add support for PandasDtype in utils.dtype
            target_dtypes: list[Dtype] = [
                np.dtype(table.schema.field(name).metadata[b'dtype'])
                for name in target_names
            ]

            ts_batch = Batch[NPArray, NPSeqLens](
                data=data_df.to_numpy(),
                seq_lens=np_seq_lens,
                feature_names=feature_names,
                dtypes=dtypes,
                targets=targets_df.to_numpy(),
                target_seq_lens=np_target_seq_lens,
                target_names=target_names,
                target_dtypes=target_dtypes,
                metadata=batch_metadata,
            )

        if as_tensor:
            ts_batch = ts_batch.unpack(pad_value=pad_value)

        if output_format and not check_output_format(output_format, category='np'):
            return ts_batch.as_format(
                output_format=output_format, pin_memory=pin_memory,
            )
        else:
            return ts_batch

    def to_parquet(
        self, filepath: str | Path,
        filesystem: fs.FileSystem | None = None,
        fs_config: dict[str, Any] | None = None,
        overwrite: bool = False,
    ) -> None:
        """Serialize the batch object to a parquet file.

        Args:
            filepath: Path to the Parquet file to create.
            filesystem: Optional [PyArrow filesystem](https://arrow.apache.org/docs/python/api/filesystems.html#filesystem-implementations) instance to write to.
            fs_config: Custom arguments to pass to [PyArrow filesystem](https://arrow.apache.org/docs/python/api/filesystems.html#filesystem-implementations) if `filesystem` is not provided.
            overwrite: If True, overwrite an existing file with the same name.

        Returns:
            None.
        """  # noqa: E501
        if not self.is_packed:
            raise ValueError('Batch must be in "packed" form to write to parquet')

        # Create a DF containing the batch data and insert an indicator column so we can
        # distinguish the data rows from the targets rows
        if isinstance(self.data, pd.DataFrame):
            df = self.data
        else:
            df = pd.DataFrame(Batch._get_numpy_view(self.data), columns=self.columns)

        df.insert(0, '__tempora_batch_data', True)

        # Add the targets to the DF
        if self.targets is not None:
            if isinstance(self.targets, pd.DataFrame):
                targets_df = self.targets
            else:
                targets_df = pd.DataFrame(
                    Batch._get_numpy_view(self.targets), columns=self.target_names,
                )
            targets_df.insert(0, '__tempora_batch_data', False)
            df = pd.concat([df, targets_df])

        # Convert the DF to a Pyarrow Table: supports read/write of metadata to parquet
        table = pa.Table.from_pandas(df)

        # Create a table metadata dict - converting np array's (seq_lens) to lists
        json_dump_np = partial(json.dumps, default=lambda x: x.tolist())
        table_metadata = {
            '__tempora_seq_lens': json_dump_np(self.seq_lens),
            '__tempora_target_seq_lens': json_dump_np(self.target_seq_lens),
            '__tempora_feature_names': json.dumps(self.feature_names),
            '__tempora_target_names': json.dumps(self.target_names),
        }
        if self.metadata:
            # Add the batch metadata if present
            table_metadata['__tempora_batch_metadata'] = json.dumps(
                self.metadata, cls=MetadataJSONEncoder,
            )

        # Store the feature / target dtypes in the column-level metadata for each field:
        # this is necessary because the Arrow types for the batch data / targets columns
        # (in the table) will not necessarily be the same as the given batch types,
        # e.g. all-null columns will have a NULL Arrow column type, and when written to
        # parquet the original dtype for the column will be lost.
        fields = [table.field('__tempora_batch_data')]
        for column, dtype in zip(self.feature_names, self.dtypes):
            field_metadata = {'dtype': str(dtype)}
            fields.append(table.field(column).with_metadata(field_metadata))
        if self.targets is not None:
            assert self.target_dtypes and self.target_names
            for column, dtype in zip(self.target_names, self.target_dtypes):
                field_metadata = {'dtype': str(dtype)}
                fields.append(table.field(column).with_metadata(field_metadata))

        # Update the table data (zero-copy) with the table and column-level metadata
        table = table.cast(pa.schema(fields, metadata=table_metadata))

        # Create a pyarrow filesystem object if one was not provided
        filepath = str(filepath)
        if filesystem is None:
            filesystem = create_pyarrow_filesystem(source=filepath, fs_config=fs_config)
            filepath = remove_uri_prefix(filepath)

        if not overwrite and filesystem.get_file_info(filepath).is_file:
            raise FileExistsError(
                f'{filepath} exists - pass "overwrite=True" to replace it',
            )

        # Create the parent directory if we're writing to local storage
        if isinstance(filesystem, fs.LocalFileSystem):
            filesystem.create_dir(str(Path(filepath).parent))

        # Write the table to a parquet file. NB: open_output_stream will automatically
        # compress the data based on the file extension.
        with filesystem.open_output_stream(filepath, compression='detect') as stream:
            pq.write_table(table, stream)

    def as_format(
        self, output_format: OutputFormat, pin_memory: bool = False,
    ) -> Batch[BatchData, SeqLens] | NoReturn:
        """Convert the batch to a different format.

        Args:
            output_format: Format of the new batch.
            pin_memory: For PyTorch tensors use page-locked CPU memory to speed up data
                transfers to GPUs.

        Returns:
            Batch object in the desired format.
        """
        format_name = get_format_name(output_format)
        if is_numeric_format(output_format) and not self.is_numeric:
            raise ValueError(
                f'Batch data must be numeric to convert to {format_name}',
            )

        # Short-circuit logic for when the target format is the current batch format
        if (
            (format_name == 'numpy' and Batch._is_numpy_batch(self))
            or (format_name == 'pandas' and Batch._is_pandas_batch(self))
            or (format_name == 'pyarrow' and Batch._is_pyarrow_batch(self))
            or (format_name == 'torch' and Batch._is_torch_batch(self))
            or (format_name == 'jax' and Batch._is_jax_batch(self))
            or (format_name == 'tensorflow' and Batch._is_tensorflow_batch(self))
        ):
            return self  # type: ignore[return-value]

        # The base format must be numpy to convert to any other format
        if not Batch._is_numpy_batch(self):
            raise ValueError(
                'Batch must be in numpy format to convert to a different format',
            )

        batch = super().as_format(output_format=output_format, pin_memory=pin_memory)
        # Need to cast here due to an issue with the IsBatch Generic Protocol...
        return cast(Batch[BatchData, SeqLens], batch)

    def select(self, features: list[str]) -> Batch[BatchData, SeqLens] | NoReturn:
        """Create a new batch object consisting of only the desired features / columns.

        Args:
            features: The features to select from the batch data.

        Returns:
            Batch object with only the selected features.
        """
        # Always ensure features is a list (just in case the user gives us a string...)
        batch = super().select(features=to_list(features))

        # For numpy batches check if the new batch is numeric and if so update the dtype
        if isinstance(batch.data, np.ndarray):
            if batch.is_numeric and not np.issubdtype(batch.data.dtype, np.number):
                if all(np.issubdtype(dt, np.integer) for dt in batch.dtypes):  # type: ignore[arg-type]  # noqa: E501
                    data = batch.data.astype(dtype=int)
                else:
                    data = batch.data.astype(dtype=float)
                batch = dc.replace(batch, data=data)
        # Need to cast here due to an issue with the IsBatch Generic Protocol...
        return cast(Batch[BatchData, SeqLens], batch)

    def select_numeric(self) -> Batch[BatchData, SeqLens] | NoReturn:
        """Create a new batch object consisting of the numeric features only."""
        numeric_features = []
        for feature_name, dt in zip(self.feature_names, self.dtypes):
            if isinstance(dt, np.dtype) and np.issubdtype(dt, np.number):
                numeric_features.append(feature_name)
        if numeric_features:
            return self.select(features=numeric_features)
        else:
            raise ValueError('Batch does not contain any numeric features/columns...')

    def _get_feature_idx(self, feature: str) -> int | NoReturn:
        if self.feature_names:
            if feature in self.feature_names:
                return self.feature_names.index(feature)
            else:
                raise ValueError(f'"{feature}" is not present in Batch feature_names')
        else:
            raise ValueError('Batch feature_names is empty...')

    def _features_to_idx(self, features: list[str]) -> list[int]:
        feature_idx: list[int] = []
        for feature in features:
            feature_idx.append(self._get_feature_idx(feature))
        return feature_idx

    def unpack(
        self, pad_value: int | float = np.nan,
    ) -> Batch[BatchData, SeqLens] | NoReturn:
        """Convert the batch to tensor form.

        Args:
            pad_value: Tensor padding value to use (for uneven sequence lengths).

        Returns:
            Batch object in tensor form.
        """
        if self.is_packed:
            # Need to cast here due to an issue with the IsBatch Generic Protocol...
            return cast(
                Batch[BatchData, SeqLens], super().unpack(pad_value=pad_value),
            )
        else:
            raise ValueError('Batch object is already unpacked / in tensor form')

    def to_tensor(
        self, pad_value: int | float = np.nan,
    ) -> Batch[BatchData, SeqLens] | NoReturn:
        """Alias for unpack()."""
        return self.unpack(pad_value=pad_value)

    @property
    def is_packed(self) -> bool:
        """True if the batch is in packed form, otherwise False."""
        return super().is_packed

    @property
    def pad_mask(self) -> PadMask | NoReturn:
        """Boolean mask of the padded data values for batches in tensor form."""
        if self.is_packed:
            raise ValueError('"pad_mask" only applicable to tensor (unpacked) form')
        else:
            return super().pad_mask

    @property
    def columns(self) -> list[str] | None:
        """Alias of the "feature_names" attribute."""
        return self.feature_names

    @property
    def feature_size(self) -> int:
        """Number of batch features / columns."""
        return self.data.shape[-1]

    @property
    def is_numeric(self) -> bool:
        """True if all batch features / columns are numeric."""
        def dtypes_numeric(dtypes: list[Dtype]) -> bool:
            return all(
                isinstance(dt, np.dtype) and np.issubdtype(dt, np.number)
                for dt in dtypes
            )

        if self.targets is not None and self.target_dtypes is not None:
            return dtypes_numeric(self.dtypes) and dtypes_numeric(self.target_dtypes)
        else:
            return dtypes_numeric(self.dtypes)

    def iter_packed(
        self, seq_lens: SeqLens | None = None,
    ) -> Iterator[slice] | NoReturn:
        """Iterate over batch data examples / sequences for a batch in packed form."""
        if not self.is_packed:
            raise ValueError('"iter_packed" can only be used with a packed Batch')
        else:
            if seq_lens is None:
                seq_lens = self.seq_lens
            start_idx = 0
            for seq_len in seq_lens:
                stop_idx = start_idx + int(seq_len)
                yield slice(start_idx, stop_idx)
                start_idx = stop_idx

    def __repr__(self) -> str:
        batch_repr = 'Batch'
        num_batch_rows_to_show = 4
        num_time_points_to_show = 6
        indent = f'{" ":>4}'

        for field in dc.fields(self):
            field_desc = field.name

            if getattr(self, field.name) is None:
                batch_repr += f'\n{indent}{field_desc}: None'
                continue

            if field.name in ('data', 'targets') and self.is_packed:
                field_desc += ' (packed)'

            if field.name in ('data', 'targets'):
                batch_repr += f'\n{indent}{field_desc}: '

                array: BatchData
                if field.name == 'data':
                    array = self.data
                    seq_lens = self.seq_lens
                else:
                    assert self.targets is not None
                    assert self.target_seq_lens is not None
                    array = self.targets
                    seq_lens = self.target_seq_lens

                # Print the data object format, shape and dtype
                batch_repr += f'{array.__class__.__name__}'
                batch_repr += f' of shape {array.shape}'
                if hasattr(array, 'dtype'):
                    dtype_repr = repr(array.dtype)
                    if dtype_repr.startswith('dtype'):
                        batch_repr += f', {dtype_repr}'
                    else:
                        batch_repr += f', dtype({dtype_repr})'

                # Determine which specific batch rows we will display
                batch_rows: list[int | None]
                batch_size = len(seq_lens)

                if batch_size <= num_batch_rows_to_show:
                    batch_rows = list(range(batch_size))
                else:
                    # Show some rows from the start and the end of the batch
                    M = int((num_batch_rows_to_show + 1) / 2)
                    batch_rows = list(range(M))  # first M rows
                    batch_rows += [None]  # None --> show an ellipsis
                    batch_rows += list(range(batch_size - M, batch_size))  # last M rows

                # Get a numpy view of the array for easy printing (zero-copy op)
                array_data = Batch._get_numpy_view(array)

                for n in batch_rows:
                    if n is None:
                        row_desc = f'{indent}...'
                        row_data = ''
                    else:
                        row_desc = f'{indent}-> segment[{n}]: '
                        row_indent = f'{" ":>{len(row_desc)}}'

                        if self.is_packed:
                            segment_slice_iter = self.iter_packed(seq_lens)
                            row_slice = next(islice(segment_slice_iter, n, None))
                            matrix = array_data[row_slice]
                        else:
                            matrix = array_data[n, slice(0, seq_lens[n])]

                        # Create a formatted string representation of the matrix
                        # that shows only the first and last N rows
                        if matrix.shape[0] <= num_time_points_to_show:
                            row_data = f'{matrix}'.replace('\n', f'\n{row_indent}')
                        else:
                            chunk_size = int(num_time_points_to_show / 2)

                            head_chunk = str(matrix[:chunk_size])
                            head_chunk = head_chunk.replace('\n', f'\n{row_indent}')
                            head_chunk = head_chunk.replace(']]', ']', 1)

                            tail_chunk = str(matrix[-chunk_size:])
                            tail_chunk = tail_chunk.replace('\n', f'\n{row_indent}')
                            tail_chunk = tail_chunk.replace('[[', f'{row_indent} [', 1)

                            row_data = f'{head_chunk}\n{row_indent} ...\n{tail_chunk}'

                    batch_repr += f'\n{row_desc}{row_data}'

            elif field.name == 'dtypes':
                batch_repr += f'\n{indent}{field_desc}: '
                # Pretty print the dtypes, e.g. show "<M8[ns]" as "datetime64[ns]"
                dtypes = ', '.join(f'{dt.name}' for dt in self.dtypes)
                batch_repr += f'[{dtypes}]'

            elif field.name == 'target_dtypes' and self.target_dtypes is not None:
                batch_repr += f'\n{indent}{field_desc}: '
                # Pretty print the dtypes, e.g. show "<M8[ns]" as "datetime64[ns]"
                dtypes = ', '.join(f'{dt.name}' for dt in self.target_dtypes)
                batch_repr += f'[{dtypes}]'

            elif field.name == 'metadata':
                # Don't show the metadata as part of the repr - can overwhelm the output
                ...

            else:
                batch_repr += f'\n{indent}{field_desc}: '
                batch_repr += repr(getattr(self, field.name))

        return batch_repr

    def equals(
        self, other: Batch[BatchData, SeqLens], *,
        check_dtype_precision: bool = True, check_metadata: bool = True,
        verbose: bool = True,
    ) -> bool:
        """Tests if the batch is equal to another batch.

        Args:
            other: The batch object to compare to.
            check_dtype_precision: Include the dtype precision in the batch comparison.
            check_metadata: Include the metadata in the batch comparison.
            verbose: Show verbose output if the two batches are not identical.

        Returns:
            True if the batches are equal, otherwise False.
        """
        # Compare the data and target values last since these take the longest to check
        names = [
            'feature_names', 'target_names', 'dtypes', 'target_dtypes',
            'seq_lens', 'target_seq_lens', 'targets', 'data',
        ]
        if check_metadata:
            names.append('metadata')

        if self.is_packed != other.is_packed:
            msg = 'Packed status differs between the two batches'
        else:
            for name in names:
                x1 = getattr(self, name)
                x2 = getattr(other, name)

                if x1 is None and x2 is None:
                    continue

                if type(x1) is not type(x2):
                    msg = f'{name} types differ: {type(x1)} vs {type(x2)}'
                    break

                if name in ('data', 'seq_lens', 'targets', 'target_seq_lens'):
                    if x1.shape != x2.shape:
                        msg = f'{name} shapes differ: {x1.shape} vs {x2.shape}'
                        break

                    np_x1 = Batch._get_numpy_view(cast(BatchData, x1))
                    np_x2 = Batch._get_numpy_view(cast(BatchData, x2))

                    if name in ('data', 'targets'):
                        if self.is_numeric:
                            # Quick check if the arrays are equal (numeric data only)
                            if np.array_equal(np_x1, np_x2, equal_nan=True):
                                continue

                        # Either arrays are numeric and differ, or they are not numeric.
                        # Loop over each feature dimension in turn and check for diffs
                        if name == 'data':
                            column_desc = 'feature'
                            column_names = self.feature_names
                        else:
                            assert self.target_names is not None
                            column_desc = 'target'
                            column_names = self.target_names

                        msg, diff_found = '', False
                        for column_idx, column_name in enumerate(column_names):
                            slice1 = np_x1[..., column_idx]
                            slice2 = np_x2[..., column_idx]

                            # For pandas format we must convert any pd.NA's to np.nan as
                            # numpy returns "elementwise comparison failed" for pd.NA
                            if isinstance(self.data, pd.DataFrame):
                                slice1[pd.isna(slice1)] = np.nan
                                slice2[pd.isna(slice2)] = np.nan

                            # When comparing the two slices elementwise we must ignore
                            # any array elements where x1 and x2 values are both NaN,
                            # since NaN != NaN
                            both_nan = pd.isna(slice1) & pd.isna(slice2)
                            diffs = ~both_nan & (slice1 != slice2)

                            if np.any(diffs):
                                diff_found = True

                                # Get the array index(es) where the first diff occurs,
                                # this will be 1-D for packed and 2-D for tensor data
                                ex_idx = np.empty(shape=slice1.ndim, dtype=int)
                                for i, diffs_idx in enumerate(np.nonzero(diffs)):
                                    ex_idx[i] = diffs_idx[0]

                                if self.is_packed:
                                    idx_desc = f'row {ex_idx[0]}'
                                else:
                                    idx_desc = f'batch {ex_idx[0]}, row {ex_idx[1]}'

                                v1, v2 = slice1[tuple(ex_idx)], slice2[tuple(ex_idx)]
                                msg += (
                                    f'{column_desc}={column_name}: {np.sum(diffs)} '
                                    f'value(s) differ, e.g. {v1} vs {v2} at '
                                    f'{idx_desc}\n'
                                )

                        if diff_found:
                            break
                    else:
                        # Check seq_lens and target_seq_lens attributes
                        if not np.array_equal(np_x1, np_x2, equal_nan=True):
                            msg = f'{name} values differ:\n{np_x1}\nvs\n{np_x2}'
                            break
                elif name == 'dtypes' and not check_dtype_precision:
                    x1 = [dtype.kind for dtype in x1]
                    x2 = [dtype.kind for dtype in x2]

                    if x1 != x2:
                        msg = f'dtype "kinds" differ:\n{x1}\nvs\n{x2}'
                        break
                else:
                    if x1 != x2:
                        msg = f'{name} differ:\n{x1}\nvs\n{x2}'
                        break
            else:
                # Didn't break so the two batches must be equal
                return True

        if verbose:
            print(msg)
        return False

    def __eq__(self, other: object) -> bool:
        if isinstance(other, Batch):
            return self.equals(other, verbose=False)
        else:
            return False

    @staticmethod
    def _get_numpy_view(data: BatchData) -> NPMatrix | NPTensor:
        # Creates a numpy view of a BatchData object - all ops below are zero copy
        # except for "pa.Table.to_pandas" which can potentially result in a mem copy
        if hasattr(data, '__class__') and 'torch' in str(data.__class__):
            return data.numpy()  # type: ignore[attr-defined, operator]
        elif hasattr(data, '__class__') and 'tensorflow' in str(data.__class__):
            # Get a read-only view of Tensorflow Tensor contents (CPU only)
            data_memview = memoryview(data)  # type: ignore[arg-type]
            return np.array(data_memview)
        elif isinstance(data, (pd.DataFrame, pa.lib.Tensor)):
            # Numpy view ensures we can index the DF/Tensor using [] notation
            return data.to_numpy()
        elif isinstance(data, pa.Table):
            # NB: passing zero_copy_only=True to to_pandas() can result in:
            # ArrowInvalid: Needed to copy 1 chunks with 0 nulls but zero_copy_only True
            return data.to_pandas(split_blocks=True).to_numpy()
        else:
            # Must be a numpy or JAX array already
            return data  # type: ignore[return-value]

    @staticmethod
    def _is_numpy_batch(
        batch: Batch[BatchData, SeqLens],
    ) -> TypeGuard[Batch[NPMatrix | NPTensor, NPSeqLens]]:
        return isinstance(batch.data, np.ndarray)

    @staticmethod
    def _is_torch_batch(
        batch: Batch[BatchData, SeqLens],
    ) -> TypeGuard[Batch[TorchMatrix | TorchTensor, TorchSeqLens]]:
        try:
            import torch
            return isinstance(batch.data, torch.Tensor)
        except ImportError:
            return False

    @staticmethod
    def _is_mutable_batch(
        batch: Batch[BatchData, SeqLens],
    ) -> TypeGuard[Batch[MutableBatchData, MutableSeqLens]]:
        return Batch._is_numpy_batch(batch) or Batch._is_torch_batch(batch)

    @staticmethod
    def _is_pandas_batch(
        batch: Batch[BatchData, SeqLens],
    ) -> TypeGuard[Batch[pd.DataFrame, NPSeqLens]]:
        return isinstance(batch.data, pd.DataFrame)

    @staticmethod
    def _is_pyarrow_batch(
        batch: Batch[BatchData, SeqLens],
    ) -> TypeGuard[Batch[pa.Table, NPSeqLens] | Batch[pa.Tensor, NPSeqLens]]:
        return isinstance(batch.data, pa.Table) or isinstance(batch.data, pa.Tensor)

    @staticmethod
    def _is_tensorflow_batch(
        batch: Batch[BatchData, SeqLens],
    ) -> TypeGuard[Batch[TFMatrix | TFTensor, TFSeqLens]]:
        try:
            import tensorflow
            return isinstance(batch.data, tensorflow.Tensor)
        except ImportError:
            return False

    @staticmethod
    def _is_jax_batch(
        batch: Batch[BatchData, SeqLens],
    ) -> TypeGuard[Batch[JaxMatrix | JaxTensor, JaxSeqLens]]:
        try:
            import jax.numpy as jnp
            return isinstance(batch.data, jnp.ndarray)
        except ImportError:
            return False


def read_batches(
    path: str | Path, *,
    filesystem: fs.FileSystem | None = None, fs_config: dict[str, Any] | None = None,
    as_tensor: bool = False, pad_value: int | float = np.nan,
    output_format: OutputFormat | None = None, pin_memory: bool = False,
) -> list[Batch[BatchData, SeqLens]]:
    """Deserialize batches from a directory of parquet files.

    Args:
        path: Directory path to the serialized batch files.
        filesystem: Optional [PyArrow filesystem](https://arrow.apache.org/docs/python/api/filesystems.html#filesystem-implementations) to read from.
        fs_config: Custom arguments to pass to [PyArrow filesystem](https://arrow.apache.org/docs/python/api/filesystems.html#filesystem-implementations) if `filesystem` is not provided.
        as_tensor: Return the batches in tensor form, otherwise (packed) matrix form.
        pad_value: Tensor padding value to use (for uneven sequence lengths).
        output_format: Format to return the batches in.
        pin_memory: For PyTorch tensors use page-locked CPU memory to speed up data transfers to GPUs.

    Returns:
        List of deserialized batch objects.
    """  # noqa: E501
    batches: list[Batch[BatchData, SeqLens]] = []

    # Create a pyarrow filesystem object if one was not provided
    if filesystem is None:
        filesystem = create_pyarrow_filesystem(source=path, fs_config=fs_config)

    # Get a list of all files under the given directory path
    base_dir = remove_uri_prefix(str(path))
    finfo_objs = filesystem.get_file_info(
        fs.FileSelector(base_dir, recursive=True),
    )
    # filesystem.get_file_info should always return list[FileInfo] if called with a
    # directory path; however if called with a filepath we must put the obj in a list
    if not isinstance(finfo_objs, list):
        finfo_objs = [finfo_objs]

    filepaths = [finfo_obj.path for finfo_obj in finfo_objs]

    for filepath in filepaths:
        try:
            batch = Batch.from_parquet(
                filepath=filepath,
                filesystem=filesystem,
                as_tensor=as_tensor,
                pad_value=pad_value,
                output_format=output_format,
                pin_memory=pin_memory,
            )
            batches.append(batch)  # type: ignore[arg-type]
        except Exception as e:
            log.warning(e.args[0])

    return batches
