# !/usr/bin/env python3
# Copyright (c) 2023 QuAIR team. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
The library of random data generation functions
"""

import math
from typing import Iterable, List, Optional, Tuple, Union

import numpy as np
import scipy
import torch
from scipy.stats import unitary_group

# TODO this is added due to channel_repr_convert, move it to intrinsic
import quairkit as qkit

from ..core import Hamiltonian, State, get_dtype, get_float_dtype, to_state
from ..core.intrinsic import _alias, _format_total_dim
from ..core.utils.linalg import _dagger

__all__ = [
    "random_pauli_str_generator",
    "random_state",
    "random_hamiltonian_generator",
    "random_hermitian",
    "random_projector",
    "random_orthogonal_projection",
    "random_density_matrix",
    "random_unitary",
    "random_unitary_hermitian",
    "random_unitary_with_hermitian_block",
    "haar_orthogonal",
    "haar_unitary",
    "haar_state_vector",
    "haar_density_operator",
    "random_channel",
    "random_clifford"
]


def random_pauli_str_generator(num_qubits: int, terms: Optional[int] = 3) -> List:
    r"""Generate a random observable in list form.

    An observable :math:`O=0.3X\otimes I\otimes I+0.5Y\otimes I\otimes Z`'s list form is
    ``[[0.3, 'x0'], [0.5, 'y0,z2']]``. Such an observable is generated by
    ``random_pauli_str_generator(3, terms=2)``.

    Args:
        num_qubits: Number of qubits.
        terms: Number of terms in the observable. Defaults to 3.

    Returns:
        The Hamiltonian of randomly generated observable.

    Examples:
        .. code-block:: python

            observable = random_pauli_str_generator(num_qubits=2, terms=2)
            print(f'The Hamiltonian of randomly generated observable is:\n{observable}')

        ::

            The Hamiltonian of randomly generated observable is:
            [[-0.6019637631250563, 'y0'], [0.12777564473712655, 'x0,z1']]
    """
    pauli_str = []
    for sublen in np.random.randint(1, high=num_qubits + 1, size=terms):
        # Tips: -1 <= coeff < 1
        coeff = np.random.rand() * 2 - 1
        ops = np.random.choice(["x", "y", "z"], size=sublen)
        pos = np.random.choice(range(num_qubits), size=sublen, replace=False)
        op_list = [ops[i] + str(pos[i]) for i in range(sublen)]
        op_list.sort(key=lambda x: int(x[1:]))
        pauli_str.append([coeff, ",".join(op_list)])
    return pauli_str


@_alias({"num_systems": "num_qubits"})
def random_state(
    num_systems: int,
    rank: Optional[int] = None,
    is_real: Optional[bool] = False,
    size: Optional[Union[int, List[int]]] = 1,
    system_dim: Union[List[int], int] = 2,
) -> State:
    r"""Generate a random quantum state.

    Args:
        num_systems: The number of qubits contained in the quantum state.
        rank: The rank of the density matrix. Defaults to ``None`` (full rank).
        is_real: If the quantum state only contains real numbers. Defaults to ``False``.
        size: Batch size. Defaults to 1.
        system_dim: Dimension of systems. Can be a list of system dimensions
            or an int representing the dimension of all systems. Defaults to the qubit case.

    Raises:
        NotImplementedError: If the backend is wrong or not implemented.

    Returns:
        The generated quantum state.

    Examples:
        .. code-block:: python

            state = random_state(num_systems=1, size=1)
            print(f'The generated quantum state is:\n{state}')

            state = random_state(num_systems=1, rank=1)
            print(f'The generated quantum state is:\n{state}')

            state = random_state(num_systems=2, system_dim=[1, 2])
            print(f'The generated quantum state is:\n{state}')

            state = random_state(num_systems=3, size=[2, 1], system_dim=[1, 2, 1])
            print(f'The generated quantum state is:\n{state}')

        ::

            The generated quantum state is:
            ---------------------------------------------------
            Backend: state_vector
            System dimension: [2]
            System sequence: [0]
            [-0.33+0.72j -0.28+0.55j]
            ---------------------------------------------------
            The generated quantum state is:
            ---------------------------------------------------
            Backend: state_vector
            System dimension: [2]
            System sequence: [0]
            [ 0.83-0.31j -0.45+0.16j]
            ---------------------------------------------------
            The generated quantum state is:
            ---------------------------------------------------
            Backend: density_matrix
            System dimension: [1, 2]
            System sequence: [0, 1]
            [[ 0.63+0.j   -0.12-0.13j]
             [-0.12+0.13j  0.37-0.j  ]]
            ---------------------------------------------------
            The generated quantum state is:
            ---------------------------------------------------
            Backend: density_matrix
            System dimension: [1, 2, 1]
            System sequence: [0, 1, 2]
            Batch size: [2, 1]
            # 0:
            [[0.43+0.j   0.24-0.37j],
             [0.24+0.37j 0.57-0.j  ]]
            # 1:
            [[ 0.88+0.j   -0.29-0.11j],
             [-0.29+0.11j  0.12-0.j  ]]
            ---------------------------------------------------
    """
    dim = _format_total_dim(num_systems, system_dim)
    size = [size] if isinstance(size, int) else list(size)
    rank = np.random.randint(1, dim + 1) if rank is None else rank
    total_size = int(np.prod(size))

    if rank == 1:
        list_state = torch.stack(
            [haar_state_vector(dim, is_real) for _ in range(total_size)]
        ).view(size + [dim, 1])
    else:
        list_state = torch.stack(
            [haar_density_operator(dim, rank, is_real) for _ in range(total_size)]
        ).view(size + [dim, dim])

    list_state = list_state if total_size > 1 else list_state.squeeze()
    return to_state(list_state, system_dim)


def random_hamiltonian_generator(
    num_qubits: int, terms: Optional[int] = 3
) -> Hamiltonian:
    r"""Generate a random Hamiltonian.

    Args:
        num_qubits: Number of qubits.
        terms: Number of terms in the Hamiltonian. Defaults to 3.

    Returns:
        The randomly generated Hamiltonian.

    Examples:
        .. code-block:: python

            Hamiltonian = random_hamiltonian_generator(3, 2)
            print(f'The randomly generated Hamiltonian is:\n{Hamiltonian}')

        ::

            The randomly generated Hamiltonian is:
            0.11801365595625102 Z0, X2
            0.9059897222160238 X0, Z1, Z2
    """
    return Hamiltonian(random_pauli_str_generator(num_qubits, terms))


def random_hermitian(num_qubits: int) -> torch.Tensor:
    r"""Randomly generate a normalized :math:`2^n \times 2^n` Hermitian matrix.

    Args:
        num_qubits: Number of qubits :math:`n`.

    Returns:
         A :math:`2^n \times 2^n` Hermitian matrix.

    Examples:
        .. code-block:: python

            Hermitian = random_hermitian(1)
            print(f'The randomly generated Hermitian matrix is:\n{Hermitian}')

        ::

            The randomly generated Hermitian matrix is:
            tensor([[0.1038+0.0000j, 0.4251+0.2414j],
                    [0.4251-0.2414j, 0.7333+0.0000j]])
    """
    assert num_qubits > 0
    n = 2**num_qubits

    mat = np.random.randn(n, n) + 1j * np.random.randn(n, n)
    for i in range(n):
        mat[i, i] = np.abs(mat[i, i])
        for j in range(i):
            mat[i, j] = np.conj(mat[j, i])

    eigval = np.linalg.eigvalsh(mat)
    max_eigval = np.max(np.abs(eigval))
    return torch.tensor(mat / max_eigval, dtype=get_dtype())


@_alias({"num_systems": "num_qubits"})
def random_projector(num_systems: int, size: int = 1,
                     system_dim: Union[List[int], int] = 2) -> torch.Tensor:
    r"""Randomly generate a :math:`d \times d` orthogonal projector.

    Args:
        num_systems: Number of systems.
        size: Number of projectors.
        system_dim: Dimension of systems. Can be a list of system dimensions 
            or an int representing the dimension of all systems. Defaults to the qubit case.

    Returns:
         A :math:`2^n \times 2^n` orthogonal projector.

    Examples:
        .. code-block:: python

            Orthogonal_projector = random_orthogonal_projection(1)
            print(f'The randomly generated Orthogonal_projector is:\n{Orthogonal_projector}')

        ::

            The randomly generated Orthogonal_projector is:
            tensor([[0.9704+1.1123e-10j, 0.1692+7.7007e-03j],
                    [0.1692-7.7008e-03j, 0.0296-1.1123e-10j]])
    """
    U = random_unitary(num_systems, size, system_dim)
    rank = np.random.randint(1, U.shape[0])
    V = U[:, :rank]
    return V @ _dagger(V)


random_orthogonal_projection = random_projector


def random_density_matrix(num_qubits: int) -> torch.Tensor:
    r"""Randomly generate an num_qubits-qubit state in density matrix form.

    Args:
        num_qubits: Number of qubits :math:`n`.

    Returns:
        A :math:`2^n \times 2^n` density matrix.

    Examples:
        .. code-block:: python

            Density_matrix = random_density_matrix(1)
            print(f'The randomly generated density matrix is:\n{Density_matrix}')

        ::

            The randomly generated density matrix is:
            tensor([[0.3380+0.0000j, 0.4579+0.1185j],
                    [0.4579-0.1185j, 0.6620+0.0000j]])
    """
    dim = 2**num_qubits
    return haar_density_operator(dim, rank=np.random.randint(1, dim + 1))


@_alias({"num_systems": "num_qubits"})
def random_unitary(
    num_systems: int,
    size: Union[int, List[int]] = 1,
    system_dim: Union[List[int], int] = 2,
) -> torch.Tensor:
    r"""Randomly generate a :math:`d \times d` unitary.

    Args:
        num_systems: Number of systems in this unitary (alias of ``num_qubits``).
        size: Batch size. Defaults to 1.
        system_dim: Dimension of systems. Can be a list of system dimensions or an int representing
            the dimension of all systems. Defaults to the qubit case.

    Returns:
         A :math:`d \times d` unitary matrix.

    Examples:
        .. code-block:: python

            unitary_matrix_1 = random_unitary(num_systems=1, system_dim=2)
            print(f'The randomly generated unitary_matrix_1 is:\n{unitary_matrix_1}')

            unitary_matrix_2 = random_unitary(num_systems=2, system_dim=[1, 2])
            print(f'The randomly generated unitary_matrix_2 is:\n{unitary_matrix_2}')

            unitary_matrix_3 = random_unitary(num_systems=1, size=[1,2])
            print(f'The randomly generated unitary_matrix_3 is:\n{unitary_matrix_3}')

        ::

            The randomly generated unitary_matrix_1 is:
            tensor([[-0.5288+0.5296j, -0.5277-0.4019j],
                    [-0.5321-0.3959j, -0.3627+0.6546j]])
            The randomly generated unitary_matrix_2 is:
            tensor([[0.6996-0.1504j, 0.3414+0.6095j],
                    [0.4954-0.4925j, -0.6315-0.3364j]])
            The randomly generated unitary_matrix_3 is:
            tensor([[[[0.3240+0.1404j, 0.4166-0.8377j],
                      [0.6068+0.7121j, -0.2804+0.2146j]],

                     [[-0.2620-0.0886j, -0.3587+0.8916j],
                      [0.5196-0.8084j, 0.2238+0.1624j]]]])
    """
    dim = _format_total_dim(num_systems, system_dim)
    size = [size] if isinstance(size, int) else list(size)
    total_size = math.prod(size)
    list_unitary = torch.stack(
        [
            torch.tensor(unitary_group.rvs(dim), dtype=get_dtype())
            for _ in range(total_size)
        ]
    ).view(size + [dim, dim])
    return list_unitary if total_size > 1 else list_unitary.squeeze()


def random_unitary_hermitian(num_qubits: int) -> torch.Tensor:
    r"""Randomly generate a :math:`2^n \times 2^n` Hermitian unitary.

    Args:
        num_qubits: Number of qubits :math:`n`.

    Returns:
         A :math:`2^n \times 2^n` Hermitian unitary matrix.

    Examples:
        .. code-block:: python

            Unitary_hermitian = random_unitary_hermitian(1)
            print(f'The randomly generated hermitian unitary is:\n{Unitary_hermitian}')

        ::

            The randomly generated hermitian unitary is:
            tensor([[0.2298+2.2018e-09j, -0.8408+4.9013e-01j],
                    [-0.8408-4.9013e-01j, -0.2298-2.2018e-09j]])
    """
    proj_mat = random_orthogonal_projection(num_qubits)
    id_mat = torch.eye(2**num_qubits)
    return (2 + 0j) * proj_mat - id_mat


def random_unitary_with_hermitian_block(
    num_qubits: int, is_unitary: bool = False
) -> torch.Tensor:
    r"""Randomly generate a unitary :math:`2^n \times 2^n` matrix that is a block encoding of a :math:`2^{n/2} \times 2^{n/2}` Hermitian matrix.

    Args:
        num_qubits: Number of qubits :math:`n`.
        is_unitary: Whether the Hermitian block is a unitary divided by 2 (for tutorial only).

    Returns:
         A :math:`2^n \times 2^n` unitary matrix whose upper-left block is Hermitian.

    Examples:
        .. code-block:: python

            unitary_matrix_1 = random_unitary_with_hermitian_block(num_qubits=2, is_unitary=False)
            print(f'The randomly generated unitary matrix 1 with hermitian block is:\n{unitary_matrix_1}')

            unitary_matrix_2 = random_unitary_with_hermitian_block(num_qubits=2, is_unitary=True)
            print(f'The randomly generated unitary matrix 2 with hermitian block is:\n{unitary_matrix_2}')

        ::

            The randomly generated unitary matrix 1 with hermitian block is:
            tensor([[ 5.7873e-01+0.0000j,  2.2460e-01+0.2711j, -1.5514e-08+0.5646j,
                      3.6316e-01-0.3009j],
                    [ 2.2460e-01-0.2711j,  7.0578e-01+0.0000j, -3.6316e-01-0.3009j,
                     -2.2489e-08+0.3944j],
                    [-1.5514e-08+0.5646j,  3.6316e-01-0.3009j,  5.7873e-01+0.0000j,
                      2.2460e-01+0.2711j],
                    [-3.6316e-01-0.3009j, -2.2489e-08+0.3944j,  2.2460e-01-0.2711j,
                      7.0578e-01+0.0000j]])
            The randomly generated unitary matrix 2 with hermitian block is:
            tensor([[-1.8185e-01-1.6847e-09j,  3.0894e-01+3.4855e-01j,
                     2.2516e-09+8.6603e-01j, -1.4451e-09-9.3500e-11j],
                    [ 3.0894e-01-3.4855e-01j,  1.8185e-01+1.6847e-09j,
                     1.7456e-09-9.3500e-11j,  2.4038e-09+8.6603e-01j],
                    [ 2.2516e-09+8.6603e-01j, -1.4451e-09-9.3500e-11j,
                    -1.8185e-01-1.6847e-09j,  3.0894e-01+3.4855e-01j],
                    [ 1.7456e-09-9.3500e-11j,  2.4038e-09+8.6603e-01j,
                     3.0894e-01-3.4855e-01j,  1.8185e-01+1.6847e-09j]])
    """
    assert num_qubits > 0

    if is_unitary:
        mat0 = random_unitary_hermitian(num_qubits - 1).detach().numpy() / 2
    else:
        mat0 = random_hermitian(num_qubits - 1).detach().numpy()
    id_mat = np.eye(2 ** (num_qubits - 1))
    mat1 = 1j * scipy.linalg.sqrtm(id_mat - np.matmul(mat0, mat0))

    mat = np.block([[mat0, mat1], [mat1, mat0]]).astype(complex)

    return torch.tensor(mat, dtype=get_dtype())


def haar_orthogonal(dim: int) -> torch.Tensor:
    r"""Randomly generate an orthogonal matrix following Haar measure,
    referenced by arXiv:math-ph/0609050v2.

    Args:
        dim: Dimension of the orthogonal matrix.

    Returns:
        A :math:`d \times d` orthogonal matrix.

    Examples:
        .. code-block:: python

            Haar_orthogonal = haar_orthogonal(2)
            print(f'The randomly generated orthogonal matrix is:\n{Haar_orthogonal}')

        ::

            The randomly generated orthogonal matrix is:
            tensor([[-0.6859+0.j,  0.7277+0.j],
                    [ 0.7277+0.j,  0.6859+0.j]])
    """
    # Step 1: sample from Ginibre ensemble
    ginibre = np.random.randn(dim, dim)
    # Step 2: perform QR decomposition of G
    mat_q, mat_r = np.linalg.qr(ginibre)
    # Step 3: make the decomposition unique
    mat_lambda = np.diag(mat_r) / abs(np.diag(mat_r))
    mat_u = mat_q @ np.diag(mat_lambda)
    return torch.tensor(mat_u, dtype=get_dtype())


def haar_unitary(dim: int) -> torch.Tensor:
    r"""Randomly generate a unitary following Haar measure,
    referenced by arXiv:math-ph/0609050v2.

    Args:
        dim: Dimension of the unitary.

    Returns:
        A :math:`d \times d` unitary.

    Examples:
        .. code-block:: python

            Haar_unitary = haar_unitary(2)
            print(f'The randomly generated unitary is:\n{Haar_unitary}')

        ::

            The randomly generated unitary is:
            tensor([[ 0.2800+0.6235j,  0.7298+0.0160j],
                    [-0.7289-0.0396j,  0.3267-0.6003j]])
    """
    # Step 1: sample from Ginibre ensemble
    ginibre = (np.random.randn(dim, dim) + 1j * np.random.randn(dim, dim)) / np.sqrt(2)
    # Step 2: perform QR decomposition of G
    mat_q, mat_r = np.linalg.qr(ginibre)
    # Step 3: make the decomposition unique
    mat_lambda = np.diag(mat_r) / np.abs(np.diag(mat_r))
    mat_u = mat_q @ np.diag(mat_lambda)
    return torch.tensor(mat_u, dtype=get_dtype())


def haar_state_vector(dim: int, is_real: Optional[bool] = False) -> torch.Tensor:
    r"""Randomly generate a state vector following Haar measure.

    Args:
        dim: Dimension of the state vector.
        is_real: Whether the vector is real. Defaults to False.

    Returns:
        A :math:`d \times 1` state vector.

    Examples:
        .. code-block:: python

            Haar_state_vector = haar_state_vector(3, is_real=True)
            print(f'The randomly generated state vector is:\n{Haar_state_vector}')

        ::

            The randomly generated state vector is:
            tensor([[0.9908+0.j],
                    [-0.1356+0.j]])
    """
    if is_real:
        # Generate a Haar random orthogonal matrix
        mat_orthog = haar_orthogonal(dim)
        # Perform orthogonal transformation on |0>
        phi = mat_orthog[:, 0]
    else:
        # Generate a Haar random unitary
        unitary = haar_unitary(dim)
        # Perform unitary transformation on |0>
        phi = unitary[:, 0]

    return phi.view([-1, 1])


def haar_density_operator(
    dim: int, rank: int, is_real: Optional[bool] = False
) -> torch.Tensor:
    r"""Randomly generate a density matrix following Haar measure.

    Args:
        dim: Dimension of the density matrix.
        rank: Rank of the density matrix.
        is_real: Whether the density matrix is real. Defaults to False.

    Returns:
        A :math:`d \times d` density matrix.

    Examples:
        .. code-block:: python

            rho1 = haar_density_operator(dim=2, rank=2)
            print(f'The randomly generated density matrix 1 is:\n{rho1}')

            rho2 = haar_density_operator(dim=2, rank=1, is_real=True)
            print(f'The randomly generated density matrix 2 is:\n{rho2}')

        ::

            The randomly generated density matrix 1 is:
            tensor([[0.8296+1.1215e-18j, -0.0430+3.5193e-01j],
                    [-0.0430-3.5193e-01j, 0.1704-1.1215e-18j]])
            The randomly generated density matrix 2 is:
            tensor([[0.6113+0.j, -0.4875+0.j],
                    [-0.4875+0.j, 0.3887+0.j]])
    """
    assert 0 < rank <= dim, "rank is an invalid number"
    if is_real:
        ginibre_matrix = np.random.randn(dim, rank)
        rho = ginibre_matrix @ ginibre_matrix.T
    else:
        ginibre_matrix = np.random.randn(dim, rank) + 1j * np.random.randn(dim, rank)
        rho = ginibre_matrix @ ginibre_matrix.conj().T
    rho = rho / np.trace(rho)
    return torch.tensor(rho, dtype=get_dtype())


@_alias({"num_systems": "num_qubits"})
def random_channel(
    num_systems: int,
    rank: int = None,
    target: str = "kraus",
    size: Optional[int] = 1,
    system_dim: Union[List[int], int] = 2,
) -> torch.Tensor:
    r"""Generate a random channel from its Stinespring representation.

    Args:
        num_systems: Number of systems.
        rank: Rank of this channel. If None, it is randomly sampled from :math:`[1, d]`.
        target: Target representation; should be ``'choi'``, ``'kraus'`` or ``'stinespring'``.
        size: Batch size. Defaults to 1.
        system_dim: Dimension of systems. Can be a list of system dimensions or an int representing
            the dimension of all systems. Defaults to the qubit case.

    Returns:
        The target representation of a random channel.
        
    Examples:
        .. code-block:: python

            channel_kraus = random_channel(num_systems=1, target="kraus", system_dim=2)
            print(f'The randomly generated kraus channel is:\n{channel_kraus}')

            channel_choi = random_channel(num_systems=1, target="choi", system_dim=2)
            print(f'The randomly generated choi channel is:\n{channel_choi}')

            channel_stinespring = random_channel(num_systems=1, target="stinespring", size=1, rank=1)
            print(f'The randomly generated stinespring channel is:\n{channel_stinespring}')

            batch_channels = random_channel(num_systems=2, size=2, target="kraus", system_dim=[1,2], rank=2)
            print(f'The randomly generated kraus channel is:\n{batch_channels}')

        ::

            The randomly generated kraus channel is:
            tensor([[[ 0.2361+0.7497j,  0.0224+0.6178j],
                     [ 0.5942-0.1706j, -0.7860+0.0085j]]])
            The randomly generated choi channel is:
            tensor([[ 0.5136+0.0000j, -0.4784-0.1446j, -0.2850-0.4106j, -0.1583-0.4886j],
                    [-0.4784+0.1446j,  0.4864+0.0000j,  0.3811+0.3023j,  0.2850+0.4106j],
                    [-0.2850+0.4106j,  0.3811-0.3023j,  0.4864+0.0000j,  0.4784+0.1446j],
                    [-0.1583+0.4886j,  0.2850-0.4106j,  0.4784-0.1446j,  0.5136+0.0000j]])
            The randomly generated stinespring channel is:
            tensor([[ 0.1652+0.5347j,  0.6310+0.5372j],
                    [ 0.4751+0.6790j, -0.5167-0.2150j]])
            The randomly generated kraus channel is:
            tensor([[[[-0.3189-0.6385j, -0.1873+0.1634j],
                      [ 0.2218+0.0440j, -0.6294+0.0947j]],

                     [[-0.0072+0.0749j, -0.0790+0.6740j],
                      [-0.5121-0.4142j, -0.0357-0.2671j]]],


                    [[[-0.1163-0.1931j,  0.2001-0.5852j],
                      [-0.5174+0.5253j,  0.1128-0.2459j]],

                     [[-0.1362+0.3280j, -0.2677+0.5444j],
                      [ 0.4328-0.3033j, -0.3862-0.1646j]]]])
    """
    target = target.lower()
    dim = _format_total_dim(num_systems, system_dim)
    rank = np.random.randint(dim) + 1 if rank is None else rank
    assert (
        1 <= rank <= dim
    ), f"rank must be positive and no larger than the dimension {dim} of the channel: received {rank}"

    list_repr = []

    for _ in range(size):

        unitary = unitary_group.rvs(rank * dim)
        stinespring_mat = torch.tensor(unitary[:, :dim], dtype=get_dtype()).reshape(
            [rank, dim, dim]
        )
        list_kraus = stinespring_mat[:rank]

        if target == "choi":
            list_repr.append(
                qkit.qinfo.channel_repr_convert(
                    list_kraus, source="kraus", target="choi"
                )
            )
        elif target == "stinespring":
            list_repr.append(
                qkit.qinfo.channel_repr_convert(
                    list_kraus, source="kraus", target="stinespring"
                )
            )
        else:
            list_repr.append(list_kraus)

    return torch.stack(list_repr) if size > 1 else list_repr[0]


def random_clifford(num_qubits: int) -> torch.Tensor:
    r"""Generate a random Clifford unitary.

    Args:
        num_qubits: The number of qubits (n).

    Returns:
        The matrix form of a random Clifford unitary.

    Reference:
        Sergey Bravyi and Dmitri Maslov, *Hadamard-free circuits expose the structure of the Clifford group*. 
        IEEE Transactions on Information Theory 67(7), 4546-4563 (2021).

    Examples:
        .. code-block:: python

            num_qubits = 1
            clifford_matrix = random_clifford(num_qubits)
            print(f'The randomly generated Clifford unitary is:\n{clifford_matrix}')

        ::

            The randomly generated Clifford unitary is:
            tensor([[ 0.7071+0.0000j, -0.7071+0.0000j],
                    [ 0.0000+0.7071j,  0.0000+0.7071j]])
    """
    # Generate a random Clifford operator (stabilizer tableaux) and the elements of the canonical form
    _, gamma_matrices, delta_matrices, hadamard_layer, permutation = __generate_random_clifford(num_qubits)

    gamma1, gamma2 = gamma_matrices[0], gamma_matrices[1]
    delta1, delta2 = delta_matrices[0].T, delta_matrices[1].T
    phase = np.random.randint(0, 2, size=2 * num_qubits)

    # Initialize the circuit
    circuit = qkit.Circuit(num_qubits)

    __apply_circuit_layers(circuit, delta2, gamma2)

    # Apply Pauli gates based on the phase vector
    x_indices, y_indices, z_indices = [], [], []
    for idx in range(num_qubits):
        if phase[idx] == 1 and phase[idx + num_qubits] == 0:
            x_indices.append(idx)
        elif phase[idx] == 0 and phase[idx + num_qubits] == 1:
            y_indices.append(idx)
        elif phase[idx] == 1 and phase[idx + num_qubits] == 1:
            z_indices.append(idx)
    if x_indices:
        circuit.x(x_indices)
    if y_indices:
        circuit.y(y_indices)
    if z_indices:
        circuit.z(z_indices)

    # Apply SWAP gates based on the permutation
    swapped_indices, swap_gate_indices = [], []
    for idx in range(num_qubits):
        if permutation[idx] == idx:
            continue
        swapped_indices.append(permutation[idx])
        if idx in swapped_indices:
            continue
        swap_gate_indices.append([idx, permutation[idx]])
    if swap_gate_indices:
        circuit.swap(swap_gate_indices)

    if hadamard_gate_indices := np.argwhere(hadamard_layer == 1).tolist():
        circuit.h(hadamard_gate_indices)

    __apply_circuit_layers(circuit, delta1, gamma1)

    # Return the Clifford unitary matrix
    return circuit.matrix


def __apply_circuit_layers(circuit, delta: Iterable[int], gamma: Iterable[int]) -> None:
    r"""Apply CNOT, CZ, and S gates to the circuit based on Delta and Gamma matrices."""
    if cnot_gate_indices := np.argwhere(np.triu(delta, k=1) == 1).tolist():
        circuit.cnot(cnot_gate_indices)
    if cz_gate_indices := np.argwhere(np.triu(gamma, k=1) == 1).tolist():
        circuit.cz(cz_gate_indices)
    if s_gate_indices := np.argwhere(np.diag(gamma) == 1).tolist():
        circuit.s(s_gate_indices)


def __generate_random_clifford(
    num_qubits: int,
) -> Tuple[np.ndarray, List[np.ndarray], List[np.ndarray], np.ndarray]:
    r"""Generate a random Clifford operator and its canonical form elements.

    TODO: bad code quality, need to eliminate the for loops.
    """
    # Constant matrices
    zero_matrix = np.zeros((num_qubits, num_qubits), dtype=int)
    zero_matrix_2n = np.zeros((2 * num_qubits, 2 * num_qubits), dtype=int)
    identity_matrix = np.eye(num_qubits, dtype=int)

    # Sample from the quantum Mallows distribution
    hadamard_layer, permutation = __sample_quantum_mallows(num_qubits)

    gamma1 = np.copy(zero_matrix)
    delta1 = np.copy(identity_matrix)
    gamma2 = np.copy(zero_matrix)
    delta2 = np.copy(identity_matrix)

    # Generate random diagonal elements for Gamma matrices
    for i in range(num_qubits):
        gamma2[i, i] = np.random.randint(2)
        if hadamard_layer[i]:
            gamma1[i, i] = np.random.randint(2)

    # Generate random elements for Gamma and Delta matrices based on canonical form constraints
    for j in range(num_qubits):
        for i in range(j + 1, num_qubits):
            b = np.random.randint(2)
            gamma2[i, j] = b
            gamma2[j, i] = b
            delta2[i, j] = np.random.randint(2)
            if hadamard_layer[i] == 1 and hadamard_layer[j] == 1:
                b = np.random.randint(2)
                gamma1[i, j] = b
                gamma1[j, i] = b
            if hadamard_layer[i] == 1 and hadamard_layer[j] == 0 and permutation[i] < permutation[j]:
                b = np.random.randint(2)
                gamma1[i, j] = b
                gamma1[j, i] = b
            if hadamard_layer[i] == 0 and hadamard_layer[j] == 1 and permutation[i] > permutation[j]:
                b = np.random.randint(2)
                gamma1[i, j] = b
                gamma1[j, i] = b
            if hadamard_layer[i] == 0 and hadamard_layer[j] == 1:
                delta1[i, j] = np.random.randint(2)
            if hadamard_layer[i] == 1 and hadamard_layer[j] == 1 and permutation[i] > permutation[j]:
                delta1[i, j] = np.random.randint(2)
            if hadamard_layer[i] == 0 and hadamard_layer[j] == 0 and permutation[i] < permutation[j]:
                delta1[i, j] = np.random.randint(2)

    # Compute stabilizer tableaux
    prod1 = np.matmul(gamma1, delta1)
    prod2 = np.matmul(gamma2, delta2)
    inv1 = np.linalg.inv(np.transpose(delta1))
    inv2 = np.linalg.inv(np.transpose(delta2))
    f1 = np.block([[delta1, zero_matrix], [prod1, inv1]])
    f2 = np.block([[delta2, zero_matrix], [prod2, inv2]])
    f1 = f1.astype(int) % 2
    f2 = f2.astype(int) % 2

    # Compute the full stabilizer tableaux
    stabilizer_tableaux = np.copy(zero_matrix_2n)
    # Apply qubit permutation to F2
    for i in range(num_qubits):
        stabilizer_tableaux[i, :] = f2[permutation[i], :]
        stabilizer_tableaux[i + num_qubits, :] = f2[permutation[i] + num_qubits, :]
    # Apply Hadamard layer
    for i in range(num_qubits):
        if hadamard_layer[i] == 1:
            stabilizer_tableaux[(i, i + num_qubits), :] = stabilizer_tableaux[(i + num_qubits, i), :]

    gamma_matrices = [gamma1, gamma2]
    delta_matrices = [delta1, delta2]

    return np.matmul(f1, stabilizer_tableaux) % 2, gamma_matrices, delta_matrices, hadamard_layer, permutation


def __sample_quantum_mallows(n: int) -> Tuple[np.ndarray, np.ndarray]:
    r"""Sample from the quantum Mallows distribution.

    Returns:
        hadamard_layer: A binary array indicating the Hadamard layer.
        permutation: A permutation array.
    """
    # Initialize Hadamard layer and permutation layer
    hadamard_layer = np.zeros(n, dtype=int)
    permutation = np.zeros(n, dtype=int)

    remaining_indices = list(range(n))

    for i in range(n):
        m = n - i  # Number of remaining indices
        r = np.random.uniform(0, 1)
        index = -1 * int(np.ceil(np.log2(r + (1 - r) * (4 ** (-1 * float(m))))))  # sample index
        hadamard_layer[i] = 1 if index < m else 0
        k = index if index < m else 2 * m - index - 1
        permutation[i] = remaining_indices[k]
        del remaining_indices[k]

    return hadamard_layer, permutation
