r"""Contain tuple-based shard generator implementations."""

from __future__ import annotations

__all__ = ["ShardTupleGenerator"]

from typing import TYPE_CHECKING, Any, TypeVar

from coola.utils import repr_indent, repr_mapping

from iden.shard import BaseShard, ShardTuple, create_shard_tuple
from iden.shard.generator.base import BaseShardGenerator, setup_shard_generator

T = TypeVar("T")

if TYPE_CHECKING:
    from pathlib import Path


class ShardTupleGenerator(BaseShardGenerator[tuple[BaseShard[T], ...]]):
    r"""Implement a shard tuple generator for creating sequences of
    shards.

    This generator creates ShardTuple instances containing an ordered
    sequence of shards, useful for organizing sequential data batches.

    Args:
        shard: The shard generator or its configuration.
        num_shards: The number of shards to generate in the
            ``ShardTuple``.
        path_uri: The path where to save the URI file.

    Example:
        ```pycon
        >>> import tempfile
        >>> import torch
        >>> from pathlib import Path
        >>> from iden.data.generator import DataGenerator
        >>> from iden.shard.generator import ShardTupleGenerator, JsonShardGenerator
        >>> with tempfile.TemporaryDirectory() as tmpdir:
        ...     generator = ShardTupleGenerator(
        ...         shard=JsonShardGenerator(
        ...             data=DataGenerator([1, 2, 3]),
        ...             path_uri=Path(tmpdir).joinpath("uri"),
        ...             path_shard=Path(tmpdir).joinpath("data"),
        ...         ),
        ...         path_uri=Path(tmpdir).joinpath("uri"),
        ...         num_shards=5,
        ...     )
        ...     generator
        ...     shard = generator.generate("shard1")
        ...     shard
        ...
        ShardTupleGenerator(
          (path_uri): PosixPath('/.../uri')
          (num_shards): 5
          (shard): JsonShardGenerator(
              (path_uri): PosixPath('/.../uri')
              (path_shard): PosixPath('/.../data')
              (data): DataGenerator(copy=False)
            )
        )
        ShardTuple(
          (uri): file:///.../uri/shard1
          (shards):
            (0): JsonShard(uri=file:///.../uri/000000001)
            (1): JsonShard(uri=file:///.../uri/000000002)
            (2): JsonShard(uri=file:///.../uri/000000003)
            (3): JsonShard(uri=file:///.../uri/000000004)
            (4): JsonShard(uri=file:///.../uri/000000005)
        )

        ```
    """

    def __init__(
        self, shard: BaseShardGenerator[T] | dict[Any, Any], num_shards: int, path_uri: Path
    ) -> None:
        self._shard = setup_shard_generator(shard)
        self._num_shards = num_shards
        self._path_uri = path_uri

    def __repr__(self) -> str:
        args = repr_indent(
            repr_mapping(
                {
                    "path_uri": self._path_uri,
                    "num_shards": self._num_shards,
                    "shard": self._shard,
                }
            )
        )
        return f"{self.__class__.__qualname__}(\n  {args}\n)"

    def equal(self, other: Any, equal_nan: bool = False) -> bool:
        if type(other) is not type(self):
            return False
        return (
            self._shard.equal(other._shard, equal_nan=equal_nan)
            and self._num_shards == other._num_shards
            and self._path_uri == other._path_uri
        )

    def generate(self, shard_id: str) -> ShardTuple[T]:
        shards = [self._shard.generate(f"{i + 1:09}") for i in range(self._num_shards)]
        return create_shard_tuple(uri=self._path_uri.joinpath(shard_id).as_uri(), shards=shards)
