from typing import Dict, List

from ._classes import (
    Labels,
    TensorBlock,
    TensorMap,
    check_isinstance,
    torch_jit_is_scripting,
)
from .slice import _slice_block


def split(
    tensor: TensorMap,
    axis: str,
    grouped_labels: List[Labels],
) -> List[TensorMap]:
    """Split a :py:class:`TensorMap` into multiple :py:class:`TensorMap`.

    The operation is based on some specified groups of indices, along either the
    "samples" or "properties" ``axis``. The number of returned :py:class:`TensorMap`s is
    equal to the number of :py:class:`Labels` objects passed in ``grouped_labels``. Each
    returned :py:class`TensorMap` will have the same keys and number of blocks at the
    input ``tensor``, but with the dimensions of the blocks reduced to only contain the
    specified indices for the corresponding group.

    For example, to split a tensor along the "samples" axis, according to the
    "structure" index, where structures 0, 6, and 7 are in the first returned
    :py:class`TensorMap`; 2, 3, and 4 in the second; and 1, 5, 8, 9, and 10 in the
    third:

    .. code-block:: python

        import metatensor

        tensor_splitted = metatensor.split(
            tensor,
            axis="samples",
            grouped_labels=[
                Labels(names=["structure"], values=np.array([[0], [6], [7]])),
                Labels(names=["structure"], values=np.array([[2], [3], [4]])),
                Labels(names=["structure"], values=np.array([[1], [5], [8], [10]])),
            ],
        )

    :param tensor: a :py:class:`TensorMap` to be split
    :param axis: a str, either "samples" or "properties", that indicates the
        :py:class:`TensorBlock` axis along which the named index (or indices) in
        ``grouped_labels`` belongs. Each :py:class:`TensorBlock` in each returned
        :py:class:`TensorMap` could have a reduced dimension along this axis, but the
        other axes will remain the same size.
    :param grouped_labels: a list of :py:class:`Labels` containing the names and values
        of the indices along the specified ``axis`` which should be in each respective
        output :py:class:`TensorMap`.

    :return: a list of:py:class:`TensorMap` that corresponds to the split input
        ``tensor``. Each tensor in the returned list contains only the named indices in
        the respective py:class:`Labels` object of ``grouped_labels``.
    """
    # Check input args
    if not torch_jit_is_scripting():
        if not check_isinstance(tensor, TensorMap):
            raise TypeError(
                f"`tensor` must be a metatensor TensorMap, not {type(tensor)}"
            )

    _check_args(tensor.block(0), axis, grouped_labels)

    all_new_blocks: Dict[int, List[TensorBlock]] = {}
    for group_i in range(len(grouped_labels)):
        empty_list: List[TensorBlock] = []
        all_new_blocks[group_i] = empty_list

    for key_index in range(len(tensor.keys)):
        key = tensor.keys.entry(key_index)
        new_blocks = _split_block(tensor[key], axis, grouped_labels)

        for group_i, new_block in enumerate(new_blocks):
            all_new_blocks[group_i].append(new_block)

    return [
        TensorMap(keys=tensor.keys, blocks=all_new_blocks[group_i])
        for group_i in range(len(grouped_labels))
    ]


def split_block(
    block: TensorBlock,
    axis: str,
    grouped_labels: List[Labels],
) -> List[TensorBlock]:
    """
    Splits an input :py:class:`TensorBlock` into multiple :py:class:`TensorBlock`
    objects based on some specified ``grouped_labels``, along either the "samples" or
    "properties" ``axis``. The number of returned :py:class:`TensorBlock`s is equal to
    the number of :py:class:`Labels` objects passed in ``grouped_labels``. Each returned
    :py:class`TensorBlock` will have the same keys and number of blocks at the input
    ``tensor``, but with the dimensions of the blocks reduced to only contain the
    specified indices for the corresponding group.

    For example, to split a block along the "samples" axis, according to the "structure"
    index, where structures 0, 6, and 7 are in the first returned :py:class`TensorMap`;
    2, 3, and 4 in the second; and 1, 5, 8, 9, and 10 in the third:

    .. code-block:: python

        import metatensor

        block_splitted = metatensor.split_block(
            block,
            axis="samples",
            grouped_labels=[
                Labels(names=["structure"], values=np.array([[0], [6]])),
                Labels(names=["structure"], values=np.array([[2], [3]])),
                Labels(names=["structure"], values=np.array([[1], [5], [10]])),
            ],
        )

    :param block: a :py:class:`TensorBlock` to be split
    :param axis: a str, either "samples" or "properties", that indicates the
        :py:class:`TensorBlock` axis along which the named index (or indices) in
        ``grouped_labels`` belongs. Each :py:class:`TensorBlock` returned could have a
        reduced dimension along this axis, but the other axes will remain the same size.
    :param grouped_labels: a list of :py:class:`Labels` containing the names and values
        of the indices along the specified ``axis`` which should be in each respective
        output :py:class:`TensorBlock`.

    :return: a list of:py:class:`TensorBlock` that corresponds to the split input
        ``block``. Each block in the returned list contains only the named indices in
        the respective py:class:`Labels` object of ``grouped_labels``.
    """
    # Check input args
    if not torch_jit_is_scripting():
        if not check_isinstance(block, TensorBlock):
            raise TypeError(
                f"`block` must be a metatensor TensorBlock, not {type(block)}"
            )

    _check_args(block, axis, grouped_labels)

    return _split_block(block, axis, grouped_labels)


def _split_block(
    block: TensorBlock,
    axis: str,
    grouped_labels: List[Labels],
) -> List[TensorBlock]:
    """
    Splits a TensorBlock into mutliple blocks, as in the public function
    :py:func:`split_block` but with no input checks. Note that the block is
    currently split into N new blocks by performing N number of slice
    operations. There may be a more efficient way of doing it, but this is not
    yet implemented.
    """
    new_blocks: List[TensorBlock] = []
    for indices in grouped_labels:
        # perform the slice either along the samples or properties axis
        new_block = _slice_block(block, axis=axis, labels=indices)
        new_blocks.append(new_block)

    return new_blocks


def _check_args(block: TensorBlock, axis: str, grouped_labels: List[Labels]):
    """
    Checks the arguments passed to :py:func:`split` and :py:func:`split_block`.
    """
    # Check types
    if not torch_jit_is_scripting():
        if not isinstance(axis, str):
            raise TypeError(f"axis must be a string, not {type(axis)}")

        if not isinstance(grouped_labels, list):
            raise TypeError(
                f"`grouped_labels` must be a list, not {type(grouped_labels)}"
            )

        for labels in grouped_labels:
            if not check_isinstance(labels, Labels):
                raise TypeError(
                    "`grouped_labels` elements must be metatensor Labels, "
                    f"not {type(labels)}"
                )

    if axis not in ["samples", "properties"]:
        raise ValueError("axis must be either 'samples' or 'properties'")

    # If passed as an empty list, return now
    if len(grouped_labels) == 0:
        return

    # Check the Labels names are equivalent for all Labels in grouped_labels
    reference_names = grouped_labels[0].names
    for labels in grouped_labels[1:]:
        if labels.names != reference_names:
            raise ValueError(
                "the dimensions names of all Labels in `grouped_labels`"
                f" must be the same, got {reference_names} and {labels.names}"
            )

    # Check the names in grouped_labels Labels are contained within the names for
    # the block
    names = block.samples.names if axis == "samples" else block.properties.names
    for name in reference_names:
        if name not in names:
            raise ValueError(
                f"the '{name}' dimension name in `grouped_labels` is not part of "
                f"the {axis} names of the input tensor"
            )
