"""
Manipulating TensorMap dimensions
=================================

Functions for manipulating dimensions of an :py:class:`metatensor.TensorMap` (i.e.
changing the columns of the :py:class:`metatensor.Labels` within).

.. autofunction:: metatensor.append_dimension

.. autofunction:: metatensor.insert_dimension

.. autofunction:: metatensor.permute_dimensions

.. autofunction:: metatensor.remove_dimension

.. autofunction:: metatensor.rename_dimension
"""
from typing import List

from ._classes import TensorBlock, TensorMap


def _check_axis(axis: str):
    if axis not in ["keys", "samples", "properties"]:
        raise ValueError(
            f"'{axis}' is not a valid axis. Choose from 'keys', 'samples' or "
            "'properties'."
        )


def append_dimension(tensor: TensorMap, axis: str, name: str, values) -> TensorMap:
    """Append a :py:class:`metatensor.Labels` dimension along the given axis.

    For ``axis=="samples"`` the new dimension is `not` appended to gradients.

    :param tensor: the input :py:class:`TensorMap`.
    :param axis: axis for which the ``name`` should be appended. Allowed are ``"keys"``,
        ``"properties"`` or ``"samples"``.
    :param name: name of the dimension be appended
    :param values: values of the dimension to be appended (``np.array`` or
        ``torch.Tensor`` according to whether ``metatensor`` or ``metatensor.torch`` is
        being used)

    :raises ValueError: if ``axis`` is a not valid value

    :return: a new :py:class:`metatensor.TensorMap` with appended labels dimension.

    Examples
    --------
    >>> import numpy as np
    >>> import metatensor
    >>> values = np.array([[1, 2], [3, 4]])
    >>> block = metatensor.block_from_array(values)
    >>> keys = metatensor.Labels(["foo"], np.array([[0]]))
    >>> tensor = metatensor.TensorMap(keys=keys, blocks=[block])
    >>> tensor
    TensorMap with 1 blocks
    keys: foo
           0
    >>> metatensor.append_dimension(
    ...     tensor,
    ...     axis="keys",
    ...     name="bar",
    ...     values=np.array([1]),
    ... )
    TensorMap with 1 blocks
    keys: foo  bar
           0    1
    """
    if axis == "keys":
        index = len(tensor.keys.names)
        return insert_dimension(tensor, axis, index, name, values)
    elif axis == "samples":
        index = len(tensor.sample_names)
        return insert_dimension(tensor, axis, index, name, values)
    elif axis == "properties":
        index = len(tensor.property_names)
        return insert_dimension(tensor, axis, index, name, values)
    else:
        raise ValueError(
            f"'{axis}' is not a valid axis. Choose from 'keys', 'samples' or "
            "'properties'."
        )


def insert_dimension(
    tensor: TensorMap,
    axis: str,
    index: int,
    name: str,
    values,
) -> TensorMap:
    """Insert a :py:class:`metatensor.Labels` dimension along the given axis before the
    given index.

    For ``axis=="samples"`` a new dimension is `not` appended to gradients.

    :param tensor: the input :py:class:`TensorMap`.
    :param axis: axis for which the ``name`` should be inserted. Allowed are ``"keys"``,
                 ``"properties"`` or ``"samples"``.
    :param index: index before the new dimension is inserted.
    :param name: the name to be inserted
    :param values: values to be inserted
        (``np.array`` or ``torch.Tensor`` according to whether ``metatensor`` or
        ``metatensor.torch`` is being used)

    :raises ValueError: if ``axis`` is a not valid value

    :return: a new :py:class:`metatensor.TensorMap` with inserted labels dimension.

    Examples
    --------
    >>> import numpy as np
    >>> import metatensor
    >>> values = np.array([[1, 2], [3, 4]])
    >>> block = metatensor.block_from_array(values)
    >>> keys = metatensor.Labels(["foo"], np.array([[0]]))
    >>> tensor = metatensor.TensorMap(keys=keys, blocks=[block])
    >>> tensor
    TensorMap with 1 blocks
    keys: foo
           0
    >>> metatensor.insert_dimension(
    ...     tensor,
    ...     axis="keys",
    ...     index=0,
    ...     name="bar",
    ...     values=np.array([1]),
    ... )
    TensorMap with 1 blocks
    keys: bar  foo
           1    0
    """
    _check_axis(axis)

    keys = tensor.keys
    if axis == "keys":
        keys = keys.insert(index=index, name=name, values=values)

    blocks: List[TensorBlock] = []
    for block in tensor.blocks():
        samples = block.samples
        properties = block.properties

        if axis == "samples":
            samples = samples.insert(index=index, name=name, values=values)
        elif axis == "properties":
            properties = properties.insert(index=index, name=name, values=values)

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

        for parameter, gradient in block.gradients():
            new_block.add_gradient(
                parameter=parameter,
                gradient=TensorBlock(
                    values=gradient.values,
                    samples=gradient.samples,
                    components=gradient.components,
                    properties=properties,
                ),
            )

        blocks.append(new_block)

    return TensorMap(keys=keys, blocks=blocks)


def permute_dimensions(
    tensor: TensorMap, axis: str, dimensions_indexes: List[int]
) -> TensorMap:
    """Permute dimensions of a :py:class:`Labels` of the given axis according to a
    :py:class:`list` of indexes.

    Values of ``dimensions_indexes`` have to be same as the indexes of
    :py:class:`Labels` but can be in a different order.

    For ``axis=="samples"`` gradients samples dimensions are not permuted.

    :param tensor: the input :py:class:`TensorMap`.
    :param axis: axis for which the ``name`` should be inserted. Allowed are ``"keys"``,
                 ``"properties"`` or ``"samples"``.
    :param dimensions_indexes: desired ordering of the dimensions

    :raises ValueError: if ``axis`` is a not valid value

    :return: a new :py:class:`metatensor.TensorMap` with the labels dimension permuted

    Examples
    --------
    >>> import numpy as np
    >>> import metatensor
    >>> values = np.array([[1, 2], [3, 4]])
    >>> block = metatensor.block_from_array(values)
    >>> keys = metatensor.Labels(["foo", "bar", "baz"], np.array([[42, 10, 3]]))
    >>> tensor = metatensor.TensorMap(keys=keys, blocks=[block])
    >>> tensor
    TensorMap with 1 blocks
    keys: foo  bar  baz
          42   10    3

    Move the last (second) dimension to the first position.

    >>> metatensor.permute_dimensions(tensor, axis="keys", dimensions_indexes=[2, 0, 1])
    TensorMap with 1 blocks
    keys: baz  foo  bar
           3   42   10
    """
    _check_axis(axis)

    keys = tensor.keys
    if axis == "keys":
        keys = keys.permute(dimensions_indexes=dimensions_indexes)

    blocks: List[TensorBlock] = []
    for block in tensor.blocks():
        samples = block.samples
        properties = block.properties

        if axis == "samples":
            samples = samples.permute(dimensions_indexes=dimensions_indexes)
        elif axis == "properties":
            properties = properties.permute(dimensions_indexes=dimensions_indexes)

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

        for parameter, gradient in block.gradients():
            new_block.add_gradient(
                parameter=parameter,
                gradient=TensorBlock(
                    values=gradient.values,
                    samples=gradient.samples,
                    components=gradient.components,
                    properties=properties,
                ),
            )

        blocks.append(new_block)

    return TensorMap(keys=keys, blocks=blocks)


def remove_dimension(tensor: TensorMap, axis: str, name: str) -> TensorMap:
    """Remove a :py:class:`metatensor.Labels` dimension along the given axis.

    Removal can only be performed if the resulting :py:class:`metatensor.Labels`
    instance will be unique.

    For ``axis=="samples"`` the dimension is not removed from gradients.

    :param tensor: the input :py:class:`TensorMap`.
    :param axis: axis for which ``name`` should be removed. Allowed are ``"keys"``,
                 ``"properties"`` or ``"samples"``.
    :param name: the :py:class:`metatensor.Labels` name to be removed

    :raises ValueError: if ``axis`` is a not valid value
    :raises ValueError: if the only dimension should be removed
    :raises ValueError: if name is not part of the axis

    :return: a new :py:class:`metatensor.TensorMap` with removed labels dimension.

    Examples
    --------
    >>> import numpy as np
    >>> import metatensor
    >>> values = np.array([[1, 2], [3, 4]])
    >>> block = metatensor.block_from_array(values)
    >>> keys = metatensor.Labels(["key", "extra"], np.array([[0, 0]]))
    >>> tensor = metatensor.TensorMap(keys=keys, blocks=[block])
    >>> tensor
    TensorMap with 1 blocks
    keys: key  extra
           0     0
    >>> metatensor.remove_dimension(tensor, axis="keys", name="extra")
    TensorMap with 1 blocks
    keys: key
           0

    Removing a dimension can only be performed if the resulting :py:class:`Labels` will
    contain unique entries.

    >>> from metatensor import MetatensorError
    >>> block = metatensor.block_from_array(values)
    >>> keys = metatensor.Labels(["key", "extra"], np.array([[0, 0], [0, 1]]))
    >>> tensor = metatensor.TensorMap(keys=keys, blocks=[block.copy(), block.copy()])
    >>> tensor
    TensorMap with 2 blocks
    keys: key  extra
           0     0
           0     1
    >>> try:
    ...     metatensor.remove_dimension(tensor, axis="keys", name="extra")
    ... except MetatensorError as e:
    ...     print(e)
    ...
    invalid parameter: can not have the same label value multiple time: [0] is already present at position 0
    """  # noqa E501
    _check_axis(axis)

    keys = tensor.keys
    if axis == "keys":
        keys = keys.remove(name=name)

    blocks: List[TensorBlock] = []
    for block in tensor.blocks():
        samples = block.samples
        properties = block.properties

        if axis == "samples":
            samples = samples.remove(name)
        elif axis == "properties":
            properties = properties.remove(name)

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

        for parameter, gradient in block.gradients():
            new_block.add_gradient(
                parameter=parameter,
                gradient=TensorBlock(
                    values=gradient.values,
                    samples=gradient.samples,
                    components=gradient.components,
                    properties=properties,
                ),
            )

        blocks.append(new_block)

    return TensorMap(keys=keys, blocks=blocks)


def rename_dimension(tensor: TensorMap, axis: str, old: str, new: str) -> TensorMap:
    """Rename a :py:class:`metatensor.Labels` dimension name for a given axis.

    :param tensor: the input :py:class:`TensorMap`.
    :param axis: axis for which the names should be appended. Allowed are ``"keys"``,
                 ``"properties"`` or ``"samples"``.
    :param old: name to be replaced
    :param new: name after the replacement

    :raises ValueError: if ``axis`` is a not valid value

    :return: a `new` :py:class:`metatensor.TensorMap` with renamed labels dimension.

    Examples
    --------
    >>> import numpy as np
    >>> import metatensor
    >>> values = np.array([[1, 2], [3, 4]])
    >>> block = metatensor.block_from_array(values)
    >>> keys = metatensor.Labels(["foo"], np.array([[0]]))
    >>> tensor = metatensor.TensorMap(keys=keys, blocks=[block])
    >>> tensor
    TensorMap with 1 blocks
    keys: foo
           0
    >>> metatensor.rename_dimension(tensor, axis="keys", old="foo", new="bar")
    TensorMap with 1 blocks
    keys: bar
           0

    """
    _check_axis(axis)

    keys = tensor.keys
    if axis == "keys":
        keys = keys.rename(old, new)

    blocks: List[TensorBlock] = []
    for block in tensor.blocks():
        samples = block.samples
        properties = block.properties

        if axis == "samples":
            samples = samples.rename(old, new)
        elif axis == "properties":
            properties = properties.rename(old, new)

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

        for parameter, gradient in block.gradients():
            gradient_samples = gradient.samples
            if axis == "samples" and old in gradient_samples.names:
                gradient_samples = gradient_samples.rename(old, new)

            new_block.add_gradient(
                parameter=parameter,
                gradient=TensorBlock(
                    values=gradient.values,
                    samples=gradient_samples,
                    components=gradient.components,
                    properties=properties,
                ),
            )

        blocks.append(new_block)

    return TensorMap(keys=keys, blocks=blocks)
