"""Determines if a channel is unitary."""

import numpy as np

from toqito.channel_ops import choi_to_kraus
from toqito.matrix_props import is_unitary as is_unitary_matrix


def is_unitary(phi: np.ndarray | list[list[np.ndarray]]) -> bool:
    r"""Given a quantum channel, determine if it is unitary.

    (Section 2.2.1: Definitions and Basic Notions Concerning Channels from
    [@watrous2018theory]).

    Let \(\mathcal{X}\) be a complex Euclidean space an let \(U \in U(\mathcal{X})\) be a
    unitary operator. Then a unitary channel is defined as:

    \[
        \Phi(X) = U X U^*.
    \]

    Examples:
        The identity channel is one example of a unitary channel:

        \[
            U =
            \begin{pmatrix}
                1 & 0 \\
                0 & 1
            \end{pmatrix}.
        \]

        We can verify this as follows:

        ```python exec="1" source="above"
        import numpy as np
        from toqito.channel_props import is_unitary

        kraus_ops = [[np.identity(2), np.identity(2)]]

        print(is_unitary(kraus_ops))
        ```

        We can also specify the input as a Choi matrix. For instance, consider the Choi matrix
        corresponding to the \(2\)-dimensional completely depolarizing channel.

        \[
            \Omega =
            \frac{1}{2}
            \begin{pmatrix}
                1 & 0 & 0 & 0 \\
                0 & 1 & 0 & 0 \\
                0 & 0 & 1 & 0 \\
                0 & 0 & 0 & 1
            \end{pmatrix}.
        \]

        We may verify that this channel is not a unitary channel.

        ```python exec="1" source="above"
        from toqito.channels import depolarizing
        from toqito.channel_props import is_unitary

        print(is_unitary(depolarizing(2)))
        ```

    Args:
        phi: The channel provided as either a Choi matrix or a list of Kraus operators.

    Returns:
        `True` if the channel is a unitary channel, and `False` otherwise.

    """
    # If the variable `phi` is provided as a ndarray, we assume this is a
    # Choi matrix.
    if isinstance(phi, np.ndarray):
        try:
            phi = choi_to_kraus(phi)
        except ValueError:
            # if we fail to obtain a Kraus representation then input/ouput spaces might be
            # non squares or their dimensions are not equal. Hence the channel is not unitary.
            return False

    # If there is a unique Kraus operator and it's a unitary matrix then the channel is unitary.
    if len(phi) != 1:
        return False

    u_mat = phi[0]
    if isinstance(phi[0], list):
        # we enter here if phi is specified as: [[U, U]] or [[U]]
        u_mat = phi[0][0]
        if len(phi[0]) > 2 or (len(phi[0]) == 2 and not np.allclose(phi[0][0], phi[0][1])):
            return False

    return is_unitary_matrix(u_mat)
