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


def slice(tensor: TensorMap, axis: str, labels: Labels) -> TensorMap:
    """
    Slice a :py:class:`TensorMap` along either the "samples" or "properties" `axis`.

    `labels` is a :py:class:`Labels` objects that specifies the
    samples/properties (respectively) names and indices that should be sliced,
    i.e. kept in the output tensor. For axis, either "samples" or "properties"
    should be specified.

    .. code-block:: python

        samples = Labels(
            names=["structure", "center"],
            values=np.array([[0, 1], [0, 6], [1, 6], [3, 16]]),  # must be a 2D-array
        )
        properties = Labels(
            names=["n"],  # radial channel
            values=np.array([[3], [4], [5]]),
        )
        sliced_tensor_samples = slice(
            tensor,
            axis="samples",
            labels=samples,
        )
        sliced_tensor_properties = slice(
            tensor,
            axis="properties",
            labels=properties,
        )

    Also note that this function will return a :py:class:`TensorMap` whose
    blocks are of equal or smaller dimensions (due to slicing) than those of the
    input. However, the returned :py:class:`TensorMap` will be returned with the
    same number of blocks and the corresponding keys as the input. If any block
    upon slicing is reduced to nothing, i.e. in the case that it has none of the
    specified `labels` along the "samples" or "properties" `axis`, an empty
    block will be returned but will still be accessible by its key. User
    warnings will be issued if any blocks are sliced to contain no values.

    For the empty blocks that may be returned, although there will be no actual
    values in its ``TensorBlock.values`` array, the shape of this array will be
    non-zero in the dimensions that haven't been sliced. This allows the slicing
    of dimensions to be tracked.

    For example, if a TensorBlock of shape (52, 1, 5) is passed, and only some
    samples are specified to be sliced but none of these appear in the input
    :py:class:`TensorBlock`, the returned block values array will be empty, but
    its shape will be (0, 1, 5) - i.e. the samples dimension has been sliced to
    zero but the components and properties dimensions remain in-tact. The same
    logic applies to any Gradient TensorBlocks the input TensorBlock may have
    associated with it.

    See the documentation for the :py:func:`slice_block` function to see how an
    individual :py:class:`TensorBlock` is sliced.

    :param tensor: the input :py:class:`TensorMap` to be sliced.
    :param axis: a :py:class:`str` indicating the axis along which slicing
        should occur. Should be either "samples" or "properties".
    :param labels: a :py:class:`Labels` object containing the names and indices
        of the "samples" or "properties" to keep in each of the sliced
        :py:class:`TensorBlock` of the output :py:class:`TensorMap`.

    :return: a :py:class:`TensorMap` that corresponds to the sliced input
        tensor.
    """
    # 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=axis, labels=labels)

    return TensorMap(
        keys=tensor.keys,
        blocks=[
            _slice_block(tensor[tensor.keys.entry(i)], axis, labels)
            for i in range(len(tensor.keys))
        ],
    )


def slice_block(block: TensorBlock, axis: str, labels: Labels) -> TensorBlock:
    """
    Slices an input :py:class:`TensorBlock` along either the "samples" or
    "properties" `axis`.  `labels` is a :py:class:`Labels` objects that specify
    the sample/property names and indices that should be sliced, i.e. kept in
    the output :py:class:`TensorBlock`.

    Example: take an input :py:class:`TensorBlock` of shape (100, 1, 6), where
    there are 100 'samples', 1 'components', and 6 'properties'. Say we want to
    slice this tensor along the samples and properties dimensions. As in the
    code-block below, we can specify, for example, 4 samples and 3 properties
    indices to keep. The returned :py:class:`TensorBlock` will have shape (4, 1,
    3).

    .. code-block:: python

        samples = Labels(
            names=["structure", "center"],
            values=np.array([[0, 1], [0, 6], [1, 6], [3, 16]]),  # must be a 2D-array
        )
        properties = Labels(
            names=["n"],  # radial channel
            values=np.array([[3], [4], [5]]),
        )
        sliced_block_samples = slice_block(
            block,
            axis="samples",
            labels=samples,
        )
        sliced_block_properties = slice_block(
            block,
            axis="properties",
            labels=properties,
        )

    For the empty blocks that may be returned, although there will be no actual
    values in its TensorBlock.values tensor, the shape of this tensor will be
    non-zero in the dimensions that haven't been sliced. This is created by
    slicing the input TensorBlock, as opposed to just returning an
    artificially-created empty one (with no shape or dimensions), and is
    intentional. It allows the slicing of dimensions to be tracked.

    For instance, if a TensorBlock of shape (52, 1, 5) is passed, and only some
    samples are specified to be sliced but none of these appear in the input
    TensorBlock, the returned TensorBlock values array will be empty, but its
    shape will be (0, 1, 5) - i.e. the samples dimension has been sliced to zero
    but the components and properties dimensions remain in-tact. The same logic
    applies to any Gradient TensorBlocks the input TensorBlock may have
    associated with it.

    :param block: the input :py:class:`TensorBlock` to be sliced.
    :param axis: a :py:class:`str` indicating the axis along which slicing
        should occur. Should be either "samples" or "properties".
    :param labels: a :py:class:`Labels` object containing the names and indices
        of the "samples" or "properties" to keep in the sliced output
        :py:class:`TensorBlock`.

    :return new_block: a :py:class:`TensorBlock` that corresponds to the sliced
        input.
    """
    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=axis, labels=labels)

    return _slice_block(
        block,
        axis=axis,
        labels=labels,
    )


def _slice_block(block: TensorBlock, axis: str, labels: Labels) -> TensorBlock:
    """
    Slices an input :py:class:`TensorBlock` along either the "samples" or
    "properties" `axis`. `labels` is :py:class:`Labels` object that specifies
    the sample/property names and indices that should be sliced, i.e. kept in
    the output :py:class:`TensorBlock`.

    :param block: the input :py:class:`TensorBlock` to be sliced.
    :param axis: a :py:class:`str` object containing `samples` or `properties`
        indicating the direction of slicing.
    :param labels: a :py:class:`Labels` object containing the names and indices
        of samples/properties to keep in the sliced output
        :py:class:`TensorBlock`.

    :return new_block: a :py:class:`TensorBlock` that corresponds to the sliced
        input.
    """

    if axis == "samples":
        # only keep the same names as `labels`
        all_samples = block.samples.view(labels.names)
        # create an arrays of bools indicating which samples indices to keep
        samples_mask = _dispatch.bool_array_like(
            [all_samples.entry(i) in labels for i in range(len(all_samples))],
            block.values,
        )
        new_values = _dispatch.mask(block.values, 0, samples_mask)
        new_samples = Labels(
            block.samples.names,
            _dispatch.mask(block.samples.values, 0, samples_mask),
        )

        new_block = TensorBlock(
            values=new_values,
            samples=new_samples,
            components=block.components,
            properties=block.properties,
        )

        # Create a map from the previous samples indexes to the new sample indexes
        # to update the gradient samples

        # sample_map contains at position old_sample the index of the
        # corresponding new sample
        sample_map = _dispatch.int_array_like(
            int_list=[-1] * len(samples_mask),
            like=samples_mask,
        )
        last = 0
        for i, picked in enumerate(samples_mask):
            if picked:
                sample_map[i] = last
                last += 1

        for parameter, gradient in block.gradients():
            if len(gradient.gradients_list()) != 0:
                raise NotImplementedError("gradients of gradients are not supported")

            # Create a samples filter for the Gradient TensorBlock
            grad_samples_mask = samples_mask[gradient.samples.column("sample")]
            new_grad_samples_values = _dispatch.mask(
                gradient.samples.values, 0, grad_samples_mask
            )

            if new_grad_samples_values.shape[0] != 0:
                # update the "sample" column of the gradient samples
                # to refer to the new samples
                new_grad_samples_values[:, 0] = sample_map[
                    new_grad_samples_values[:, 0]
                ]

                new_grad_samples = Labels(
                    names=gradient.samples.names,
                    values=new_grad_samples_values,
                )
            else:
                new_grad_samples = Labels(
                    names=gradient.samples.names,
                    values=_dispatch.empty_like(
                        gradient.samples.values, [0, gradient.samples.values.shape[1]]
                    ),
                )

            new_grad_values = _dispatch.mask(gradient.values, 0, grad_samples_mask)
            # Add sliced gradient to the TensorBlock
            new_block.add_gradient(
                parameter=parameter,
                gradient=TensorBlock(
                    values=new_grad_values,
                    samples=new_grad_samples,
                    components=gradient.components,
                    properties=new_block.properties,
                ),
            )

    else:
        assert axis == "properties"
        # only keep the same names as `labels`
        all_properties = block.properties.view(list(labels.names))
        # create an arrays of bools indicating which samples indices to keep
        properties_mask = _dispatch.bool_array_like(
            [all_properties.entry(i) in labels for i in range(len(all_properties))],
            block.values,
        )
        new_values = _dispatch.mask(
            block.values, len(block.values.shape) - 1, properties_mask
        )
        new_properties = Labels(
            block.properties.names,
            _dispatch.mask(block.properties.values, 0, properties_mask),
        )

        new_block = TensorBlock(
            values=new_values,
            samples=block.samples,
            components=block.components,
            properties=new_properties,
        )

        # Slice each Gradient TensorBlock and add to the new_block.
        for parameter, gradient in block.gradients():
            if len(gradient.gradients_list()) != 0:
                raise NotImplementedError("gradients of gradients are not supported")

            assert axis == "properties"
            new_grad_values = _dispatch.mask(
                gradient.values, len(gradient.values.shape) - 1, properties_mask
            )
            new_grad_samples = gradient.samples

            # Add sliced gradient to the TensorBlock
            new_block.add_gradient(
                parameter=parameter,
                gradient=TensorBlock(
                    values=new_grad_values,
                    samples=new_grad_samples,
                    components=gradient.components,
                    properties=new_block.properties,
                ),
            )

    return new_block


def _check_args(
    block: TensorBlock,
    axis: str,
    labels: Labels,
):
    """
    Checks the arguments passed to :py:func:`slice` and :py:func:`slice_block`.
    """
    # check axis
    if axis not in ["samples", "properties"]:
        raise ValueError(
            f"``axis``: {axis} is not known as a slicing axis. Please use"
            "'samples' or 'properties'"
        )

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

    if axis == "samples":
        s_names = block.samples.names
        for name in labels.names:
            if name not in s_names:
                raise ValueError(
                    f"invalid sample name '{name}' which is not part of the input"
                )
    else:
        assert axis == "properties"
        p_names = block.properties.names
        for name in labels.names:
            if name not in p_names:
                raise ValueError(
                    f"invalid property name '{name}' which is not part of the input"
                )
