"""Computes a list of Kraus operators from the Choi matrix."""

import numpy as np

from toqito.channel_props.channel_dim import channel_dim
from toqito.matrix_ops import unvec
from toqito.matrix_props import is_hermitian, is_positive_semidefinite


def choi_to_kraus(
    choi_mat: np.ndarray, tol: float = 1e-9, dim: int | list[int] | np.ndarray | None = None
) -> list[np.ndarray] | list[list[np.ndarray]]:
    r"""Compute a list of Kraus operators from the Choi matrix from [@rigetti2022forest].

    Note that unlike the Choi or natural representation of operators, the Kraus representation is
    *not* unique.

    If the input channel maps \(M_{r,c}\) to \(M_{x,y}\) then `dim` should be the
    list `[[r,x], [c,y]]`. If it maps \(M_m\) to \(M_n\), then `dim` can simply
    be the vector `[m,n]`.

    For completely positive maps the output is a single flat list of numpy arrays since the left and
    right Kraus maps are the same.

    This function has been adapted from [@rigetti2022forest] and QETLAB [@qetlablink].

    Examples:
        Consider taking the Kraus operators of the Choi matrix that characterizes the "swap operator"
        defined as

        \[
            \begin{pmatrix}
                1 & 0 & 0 & 0 \\
                0 & 0 & 1 & 0 \\
                0 & 1 & 0 & 0 \\
                0 & 0 & 0 & 1
            \end{pmatrix}
        \]

        The corresponding Kraus operators of the swap operator are given as follows,

        \[
            \begin{equation}
            \big[
                \frac{1}{\sqrt{2}} \begin{pmatrix} 0 & 1 \\ -1 & 0 \end{pmatrix},
                \frac{1}{\sqrt{2}} \begin{pmatrix} 0 & -1 \\ 1 & 0 \end{pmatrix}
            \big],
            \big[
                \frac{1}{\sqrt{2}} \begin{pmatrix} 0 & 1 \\ 1 & 0 \end{pmatrix},
                \frac{1}{\sqrt{2}} \begin{pmatrix} 0 & 1 \\ 1 & 0 \end{pmatrix}
            \big],
            \big[
                \begin{pmatrix} 1 & 0 \\ 0 & 0 \end{pmatrix},
                \begin{pmatrix} 1 & 0 \\ 0 & 0 \end{pmatrix}
            \big],
            \big[
                \begin{pmatrix} 0 & 0 \\ 0 & 1 \end{pmatrix},
                \begin{pmatrix} 0 & 0 \\ 0 & 1 \end{pmatrix}
            \big]
            \end{equation}
        \]

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

        ```python exec="1" source="above"
        import numpy as np
        from toqito.channel_ops import choi_to_kraus
        choi_mat = np.array([[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]])
        kraus_ops = choi_to_kraus(choi_mat)
        for i, pair in enumerate(kraus_ops):
           print(f"\nKraus Pair {i+1}:")
           for j, op in enumerate(pair):
               print(f"  Operator {j+1}:\n{np.array_str(op, precision=4, suppress_small=True)}")
        ```

        !!! See Also
            [kraus_to_choi][toqito.channel_ops.kraus_to_choi.kraus_to_choi]

    Args:
        choi_mat: A Choi matrix
        tol: optional threshold parameter for eigenvalues/kraus ops to be discarded
        dim: A scalar, vector or matrix containing the input and output dimensions of Choi matrix.

    Returns:
        List of Kraus operators

    """
    d_in, d_out, _ = channel_dim(choi_mat, dim=dim, compute_env_dim=False)
    if is_hermitian(choi_mat):
        eigvals, v_mat = np.linalg.eigh(choi_mat)
        kraus_0 = [
            np.sqrt(abs(eigval)) * unvec(evec, shape=(d_out[0], d_in[0]))
            for eigval, evec in zip(eigvals, v_mat.T)
            if abs(eigval) > tol
        ]

        if is_positive_semidefinite(choi_mat):
            return kraus_0

        kraus_1 = [
            np.sign(eigval) * k_mat for eigval, k_mat in zip(filter(lambda eigval: abs(eigval) > tol, eigvals), kraus_0)
        ]
    else:
        u_mat, singular_values, vh_mat = np.linalg.svd(choi_mat, full_matrices=False)
        kraus_0 = [
            np.sqrt(s_val) * unvec(evec, shape=(d_out[0], d_in[0]))
            for s_val, evec in zip(singular_values, u_mat.T)
            if abs(s_val) > tol
        ]

        kraus_1 = [
            np.sqrt(s_val) * unvec(evec.conj(), shape=(d_out[1], d_in[1]))
            for s_val, evec in zip(singular_values, vh_mat)
            if abs(s_val) > tol
        ]

    return [[ka, kb] for ka, kb in zip(kraus_0, kraus_1)]
