"""Tensor product operation calculates the tensor product between vectors or matrices."""

import numpy as np


def tensor(*args: np.ndarray | int | list[np.ndarray]) -> np.ndarray | None:
    r"""Compute the Kronecker tensor product [@wikipediatensor].

    Tensor two matrices or vectors together using the standard Kronecker
    operation provided from numpy.

    Given two matrices \(A\) and \(B\), computes \(A \otimes B\).
    The same concept also applies to two vectors \(v\) and \(w\) which
    computes \(v \otimes w\).

    One may also compute the tensor product one matrix \(n\) times with itself.

    For a matrix, \(A\) and an integer \(n\), the result of this
    function computes \(A^{\otimes n}\).

    Similarly for a vector \(v\) and an integer \(n\), the result of
    of this function computes \(v^{\otimes n}\).

    One may also perform the tensor product on a list of matrices.

    Given a list of \(n\) matrices \(A_1, A_2, \ldots, A_n\) the result
    of this function computes

    \[
        A_1 \otimes A_2 \otimes \cdots \otimes A_n.
    \]

    Similarly, for a list of \(n\) vectors \(v_1, v_2, \ldots, v_n\),
    the result of this function computes

    \[
        v_1 \otimes v_2 \otimes \cdots \otimes v_n.
    \]

    Examples:
        Tensor product two matrices or vectors

        Consider the following ket vector

        \[
            e_0 = \left[1, 0 \right]^{\text{T}}.
        \]

        Computing the following tensor product

        \[
        e_0 \otimes e_0 = [1, 0, 0, 0]^{\text{T}}.
        \]

        This can be accomplished in `|toqito⟩` as follows.

        ```python exec="1" source="above"
        from toqito.states import basis
        from toqito.matrix_ops import tensor

        e_0 = basis(2, 0)

        print(tensor(e_0, e_0))
        ```


        Tensor product one matrix \(n\) times with itself.

        We may also tensor some element with itself some integer number of times.
        For instance we can compute

        \[
            e_0^{\otimes 3} = \left[1, 0, 0, 0, 0, 0, 0, 0 \right]^{\text{T}}
        \]

        in `|toqito⟩` as follows.

        ```python exec="1" source="above"
        from toqito.states import basis
        from toqito.matrix_ops import tensor

        e_0 = basis(2, 0)

        print(tensor(e_0, 3))
        ```

        Perform the tensor product on a list of vectors or matrices.

        If we wish to compute the tensor product against more than two matrices or
        vectors, we can feed them in as a `list`. For instance, if we wish to
        compute \(e_0 \otimes e_1 \otimes e_0\), we can do
        so as follows.

        ```python exec="1" source="above"
        from toqito.states import basis
        from toqito.matrix_ops import tensor

        e_0, e_1 = basis(2, 0), basis(2, 1)

        print(tensor([e_0, e_1, e_0]))
        ```

    Raises:
        ValueError: Input must be a vector or matrix.

    Args:
        args: Input to the tensor function is expected to be either:
            - list[np.ndarray]: List of numpy matrices,
            - np.ndarray, ... , np.ndarray: An arbitrary number of numpy arrays,
            - np.ndarray, int: A numpy array and an integer.

    Returns:
        The computed tensor product.

    """

    def fast_exp(matrix: np.ndarray, q: int) -> np.ndarray:
        """Efficient exponentiation by squaring."""
        if q == 1:
            return matrix
        tmp = fast_exp(matrix, q >> 1)
        tmp = np.kron(tmp, tmp)
        if q & 1:  # If q is odd
            tmp = np.kron(matrix, tmp)
        return tmp

    result = None

    # Input is provided as a list of numpy matrices.
    if (len(args) == 1 and isinstance(args[0], list)) or (len(args) == 1 and isinstance(args[0], np.ndarray)):
        if len(args[0]) == 1:
            return args[0][0]
        if len(args[0]) == 2:
            return np.kron(args[0][0], args[0][1])
        if len(args[0]) >= 3:
            result = args[0][0]
            for i in range(1, len(args[0])):
                result = np.kron(result, args[0][i])
        return result

    # Tensor product one matrix `n` times with itself.
    if len(args) == 2 and isinstance(args[1], int):
        num_tensor = args[1]
        if num_tensor == 0:
            return np.eye(1, dtype=args[0].dtype)
        if num_tensor == 1:
            return args[0]
        return fast_exp(args[0], num_tensor)

    # Tensor product between two or more matrices.
    if len(args) == 2:
        return np.kron(args[0], args[1])
    if len(args) >= 3:
        result = args[0]
        for i in range(1, len(args)):
            result = np.kron(result, args[i])
        return result

    raise ValueError("The `tensor` function must take either a matrix or vector.")
