"""Utility functions used by mGST modules."""
# ruff: noqa: N802, N803, N806

import os
import random
import warnings

from iqm.benchmarks.compressive_gst.mgst.low_level_jit import Mp_norm_lower, MVE_lower, contract, local_basis
import numpy as np
import numpy.linalg as la
from scipy.linalg import expm, qr


def transp(dim1: int, dim2: int) -> np.ndarray:
    """Superoperator of a map that performs the transpose operation.

    Args:
        dim1: First dimension of the matrix to be transposed
        dim2: Second dimension of the matrix to be transposed

    Returns:
        T: Transpose superoperator that turns a dim1xdim2 matrix into a dim2xdim1 matrix

    Notes:
        The projection is done with respect to the canonical metrik.

    """
    id1 = np.eye(dim1)
    id2 = np.eye(dim2)
    T = np.einsum("il,jk->ijkl", id2, id1).reshape(dim1 * dim2, dim1 * dim2)
    return T


def randpsd(n: int, normalized: str = "True") -> np.ndarray:
    """Generate a random positive semidefinite square matrix.

    Args:
        n: Number of matrix entries
        normalized: Controls if the output is trace normalized

    Returns:
        mat: Random positive definite square matrix

    Notes:
        Eigenvalues are drawn uniformly from the interval [0,1);
        The basis is generated by diagonalizing a random hermitian matrix (see randHerm function).

    """
    dim = int(np.sqrt(n))
    H = randHerm(dim)
    U, _, _ = np.linalg.svd(H)
    evals = np.random.random_sample(dim)
    if normalized == "True":
        evals /= np.sum(evals)
    mat = np.dot(np.dot(U, np.diag(evals)), U.T.conj())
    mat = mat.reshape(-1)
    return mat


def randvec(n: int) -> np.ndarray:
    """Generate vector with real and imaginary part drawn from the normal distribution.

    Args:
        n: Number of elements for the random vector

    Returns:
        g: Length n vector with complex entries whose real and imaginary part is independently drawn
            from the normal distribution with mean 0 and variance 1.

    """
    # randn(n) produces a random vector of length n with mean 0 and variance 1
    g = np.random.randn(n) + 1j * np.random.randn(n)
    g = g / np.linalg.norm(g)
    return g


def randHerm(n: int) -> np.ndarray:
    """Generate random square hermitian matrix.

    Args:
        n: Matrix dimension

    Returns:
        G: Random hermitian matrix normalized in spectral norm

    Notes:
        First a matrix with random complex entries is generated (see randvec function).
        This matrix is then projected onto the space of hermitian
        matrices and normalized in spectral norm.

    """
    G = randvec(n * n).reshape(n, n)
    G = (G + G.T.conj()) / 2
    # ord=2 gives the spectral norm
    G = G / np.linalg.norm(G, ord=2)
    return G


def randHermGS(d: int, r: int) -> np.ndarray:
    """Generates random set of operators that are hermiticity preserving.

    Args:
        d: Number of gates
        r: Superoperator dimension of the gates given by the square of the physical dimension

    Returns:
        X: Array where random hermiticity preserving operators are stacked along the first axis.

    Notes:
        The function randHerm is used to generate random hermitian Choi matrices,
        whose indies are then rearanged to obtain random hermiticity preserving superoperators.

    """
    dim = int(np.sqrt(r))
    X = np.zeros((r, d, r), dtype="complex")
    for i in range(d):
        H = randHerm(r).reshape(dim, dim, dim, dim)
        H = np.einsum("ijkl->jlik", H)
        X[:, i, :] = H.reshape(r, r)
    return X


def randU(n: int, a: float = 1) -> np.ndarray:
    """Generates random unitary from a random hermitian generator.

    Args:
        n: Matrix dimension of the unitary
        a: Parameter to control the norm of the hermitian generator

    Returns:
        U: Matrix exponential of random hermitian matrix times the imaginary unit.

    """
    return expm(1j * a * randHerm(n)).astype(np.complex128)


def randU_Haar(n: int) -> np.ndarray:
    """Return a Haar distributed random unitary.

    Args:
        n: Matrix dimension of the unitary

    Returns:
        U: Random unitary matrix distributed according to the Haar measure.

    """
    Z = np.random.randn(n, n) + 1.0j * np.random.randn(n, n)
    [Q, R] = qr(Z)
    D = np.diag(np.diagonal(R) / np.abs(np.diagonal(R)))
    return np.dot(Q, D)


def randKrausSet(d: int, r: int, rK: int, a: float = 1) -> np.ndarray:
    """Generates random set of Kraus operators.

    Args:
        d: Number of gates
        r: Superoperator dimension of the gates given by the square of the physical dimension
        rK: Number of Kraus operators per gate ("Kraus rank")
        a: Parameter to control the norm of the hermitian generator and thereby how far the gates are from the identity

    Returns:
        K: Each subarray along the first axis contains a set of Kraus operators.
            The second axis enumerates Kraus operators for a gate specified by the first axis.

    Notes:
        Let pdim be the physical dimension. Then a set of Kraus operators is generated by taking the first pdim
            columns of a random unitary of size rK*pdim. The random unitary is generated from a random hermitian matrix.

    """
    pdim = int(np.sqrt(r))
    K = np.zeros((d, rK, pdim, pdim)).astype(np.complex128)
    for i in range(d):
        K[i, :, :, :] += randU(pdim * rK, a)[:, :pdim].reshape(rK, pdim, pdim)
    return K


def randKrausSet_Haar(d: int, r: int, rK: int) -> np.ndarray:
    """Generates random set of Kraus operators.

    Args:
        d: Number of gates
        r: Superoperator dimension of the gates given by the square of the physical dimension
        rK: Number of Kraus operators per gate ("Kraus rank")

    Returns:
        K: Each subarray along the first axis contains a set of Kraus operators.
            The second axis enumerates Kraus operators for a gate specified by the first axis.

    Notes:
        Let pdim be the physical dimension. Then a set of Kraus operators is generated by taking the first
            pdim columns of a random unitary of size rK*pdim. The random unitary is generated according to
            the Haar measure.

    """
    pdim = int(np.sqrt(r))
    K = np.zeros((d, rK, pdim, pdim)).astype(np.complex128)
    for i in range(d):
        K[i, :, :, :] += randU_Haar(pdim * rK)[:, :pdim].reshape(rK, pdim, pdim)
    return K


def random_gs(d: int, r: int, rK: int, n_povm: int) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Generates a random gate using the Gaussian unitary ensemble, initial state and POVM.

    Args:
        d: Number of gates
        r: Superoperator dimension of the gates given by the square of the physical dimension
        rK: Number of Kraus operators per gate ("Kraus rank")
        n_povm: Number of POVM-Elements

    Returns:
        K: Each subarray along the first axis contains a set of Kraus operators.
            The second axis enumerates Kraus operators for a gate specified by the first axis.
        X: Array where random CPT superoperators are stacked along the first axis.
        E: Randomly generated POVM elements stacked along the first axis.
        rho: Random initial state vector

    Notes:
        The Kraus operators are generated from random unitaries, see function randKrausSet

    """
    K = randKrausSet(d, r, rK).copy()
    X = np.einsum("ijkl,ijnm -> iknlm", K, K.conj()).reshape((d, r, r))
    rho = randpsd(r).copy()
    A = randKrausSet(1, r, n_povm)[0].conj()
    E = np.array([(A[i].T.conj() @ A[i]).reshape(-1) for i in range(n_povm)]).copy()
    return K, X, E, rho


def random_gs_Haar(d: int, r: int, rK: int, n_povm: int) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Generates a random gate set with gates from Haar random unitaries, initial state and POVM.

    Args:
        d: Number of gates
        r: Superoperator dimension of the gates given by the square of the physical dimension
        rK: Number of Kraus operators per gate ("Kraus rank")
        n_povm: Number of POVM-Elements

    Returns:
        K: Each subarray along the first axis contains a set of Kraus operators.
            The second axis enumerates Kraus operators for a gate specified by the first axis.
        X: Array where random CPT superoperators are stacked along the first axis.

    Notes:
        The Kraus operators are generated from Haar random unitaries, see function randKrausSet_Haar

    """
    K = randKrausSet_Haar(d, r, rK).copy()
    X = np.einsum("ijkl,ijnm -> iknlm", K, K.conj()).reshape((d, r, r))
    rho = randpsd(r).copy()
    A = randKrausSet_Haar(1, r, n_povm)[0].conj()
    E = np.array([(A[i].T.conj() @ A[i]).reshape(-1) for i in range(n_povm)]).copy()
    return K, X, E, rho


def perturbed_target_init(X_target: np.ndarray, rK: int) -> np.ndarray:
    """Generates a small random noise gate around the identity and applies it to the target gate.

    The reason for using this gate as an initialization as opposed the target gate itself, is
    that the non-dominant Kraus operators we start with are now not zero, but small random matrices.
    Observations show that with this start, the non-dominant Kraus operators converge faster.

    Args:
        X_target: Current gate estimate
        rK: Number of Kraus operators per gate ("Kraus rank") for the initialization

    Returns:
        K_init: Each subarray along the first axis contains a set of Kraus operators.
            The second axis enumerates Kraus operators for a gate specified by the first axis.

    """
    d, r, _ = X_target.shape
    pdim = int(np.sqrt(r))
    K_perturb = randKrausSet(d, r, rK, a=0.05)
    X_perturb = np.einsum("ijkl,ijnm -> iknlm", K_perturb, K_perturb.conj()).reshape((d, r, r))
    X_init = np.einsum("ikl,ilm ->ikm", X_perturb, X_target)
    K_init = Kraus_rep(X_init, d, pdim, rK)
    return K_init


def basis(size: int, index: int) -> np.ndarray:
    """Creates standard basis vectors.

    Args:
        size: Vector space dimension
        index: Index of basis vector

    Returns:
        vec: Vector with entry 1 at position given by index and zeros elsewhere

    """
    vec = np.zeros(size)
    vec[index] = 1.0
    return vec


def depol(pdim: int, p: float) -> np.ndarray:
    """Kraus representation of depolarizing channel.

    Args:
        pdim: Physical dimension
        p: Error probability

    Returns:
        K_depol: Each entry along the first axis contains a set of Kraus operators. The second axis enumerates
            Kraus operators for a gate specified by the first axis.

    Notes:
        The depolarizing channel is defined as L(rho) = (1-p)*rho + p/pdim*Id.

    """
    phi_plus = np.sum([np.kron(basis(pdim, i), basis(pdim, i)) for i in range(pdim)], axis=0)
    choi_state = p / pdim * np.eye(pdim**2) + (1 - p) * np.kron(phi_plus, phi_plus.reshape(pdim**2, 1))
    K_depol = la.cholesky(choi_state)
    return K_depol.reshape(pdim, pdim, pdim**2).swapaxes(0, 2)


def batch(y: np.ndarray, J: np.ndarray, bsize: int) -> tuple[np.ndarray, np.ndarray]:
    """Returns random batch of sequences and corresponding measurements.

    Args:
        y: 2D array of measurement outcomes for sequences in J; Each column contains the outcome probabilities for
            a fixed sequence
        J: 2D array where each row contains the gate indices of a gate sequence
        bsize: Size of the batch (number of sequences)

    Returns:
        y_b: Randomly subsampled columns of y
        J_b: Randomly subsampled rows of J in accordance the the columns selected in y_b

    """
    if y.shape[1] <= bsize:
        return y, J
    if bsize < 1:  # if batch size is given as ratio
        bsize = int(bsize * len(J) // 1)
    batchmask = np.array([1] * bsize + [0] * (len(J) - bsize))
    np.random.shuffle(batchmask)
    J_b = J[batchmask == 1]
    y_b = y[:, batchmask == 1]
    return y_b, J_b


def F_avg_X(X: np.ndarray, K: np.ndarray) -> tuple[float, np.ndarray]:
    r"""Returns the average gate fidelity between a superoperator and a set of Kraus operators.

    Args:
        X: CPT superoperator of size (pysical dimension ** 2) x (pysical dimension ** 2) as 2D array.
        K: Array of Kraus operators with size (Kraus rank) x (pysical dimension) x (pysical dimension) as 3D array.

    Returns:
        Average gate fidelity.

    Notes:
        No gauge optimization is involved; average gate fidelity is gauge dependent.

    """
    pdim = K.shape[2]
    d = K.shape[0]
    Fid_list = []
    for k in range(d):
        choi_inner_prod = np.einsum("imjl,pml,pij", X[k].reshape(pdim, pdim, pdim, pdim), K[k], K[k].conj())

        unitality_term = np.einsum("imkk,pml,pil", X[k].reshape(pdim, pdim, pdim, pdim), K[k], K[k].conj())
        Fid = (choi_inner_prod + unitality_term) / pdim / (pdim + 1)
        Fid_list.append(Fid)
    return np.average(np.real(Fid_list)), np.real(Fid_list)


def MVE(
    gates_true: list[np.ndarray],
    gates_estimate: list[np.ndarray],
    d: int,
    length: int,
    n_povm: int,
    samples: int = 10000,
) -> tuple[float, float]:
    """Mean varation error between the outputs of two gate sets on random sequences.

    Args:
        gates_true: Target gates, POVM and initial state
        gates_estimate: Current gate, POVM and initial state estimate
        d: Number of different gates in the gate set
        length: Length of the test sequences
        n_povm: Number of POVM elements
        samples: Number of random gate sequences over which the mean variation error is computed

    Returns:
        MVE: Mean varaition error

    Notes:
        Sequences are drawn without replacement from initally d**l possibilities.
        For each sequence the total variation error of the two probability distribution
        over the POVM elements is computed. Afterwards the meean over these total
        variation errors is returned.

    """
    X_true, E_true, rho_true = gates_true
    X, E, rho = gates_estimate
    if samples == "all" or np.log(samples) / np.log(d) > length:
        J = np.random.randint(0, d, length * d**length).reshape(d**length, length)
    else:
        J = np.random.randint(0, d, length * samples).reshape(samples, length)
    return MVE_lower(X_true, E_true, rho_true, X, E, rho, J, n_povm)


def Mp_norm(
    gates_true: list[np.ndarray],
    gates_estimate: list[np.ndarray],
    d: int,
    length: int,
    p: int,
    samples: int = 10000,
) -> tuple[float, float]:
    """Mean of the p-norm deviation between the outputs of two gate sets on random sequences.

    Args:
        gates_true: Target gates, POVM and initial state
        gates_estimate: Current gate, POVM and initial state estimate
        d: Number of different gates in the gate set
        length: Length of the test sequences
        p: Defines the l_p - norm that is used to compare probability distributions
        samples: Number of random gate sequences over which the mean variation error is computed

    Returns:
        MPE: Mean l_p - norm error

    Notes:
        Sequences are drawn without replacement from initally d**l possibilities.
        For each sequence the l_p - norm error of the two probability distribution
        over the POVM elements is computed. Afterward the meean over these total
        variation errors is returned.

    """
    X_true, E_true, rho_true = gates_true
    X, E, rho = gates_estimate
    if samples == "all" or np.log(samples) / np.log(d) > length:
        J = np.random.randint(0, d, length * d**length).reshape(d**length, length)
    else:
        J = np.random.randint(0, d, length * samples).reshape(samples, length)
    return Mp_norm_lower(X_true, E_true, rho_true, X, E, rho, J, p)


def Kraus_rep(X: np.ndarray, d: int, pdim: int, rK: int) -> np.ndarray:
    """Compute the Kraus representations for all gates in the gate set.

    Args:
        X: Current gate estimate
        d: Number of gates
        pdim: Physical dimension
        rK: Target Kraus rank

    Returns:
        K: Each subarray along the first axis contains a set of Kraus operators.
            The second axis enumerates Kraus operators for a gate specified by the first axis.

    Notes:
        The Kraus representation is obtained from a singular value decomposition
        of the Choi matrix. If parameter rK is smaller than the true rank of the
        Choi matrix, a rank rK approximation is used. Approximations are only CPT
        in the special case where the original gate was already of rank <= rK.

    """
    X_choi = X.reshape(d, pdim, pdim, pdim, pdim)
    X_choi = np.einsum("ijklm->iljmk", X_choi).reshape((d, pdim**2, pdim**2))
    K = np.zeros((d, rK, pdim, pdim)).astype(np.complex128)
    psd_tolerance = 1e-12
    for i in range(d):
        w, v = la.eigh(X_choi[i])
        if np.min(w) < -psd_tolerance:
            raise ValueError(f"Choi Matrix is not positive definite within tolerance {psd_tolerance}")
        K[i] = np.einsum(
            "ijk->kji",
            (v[:, -rK:] @ np.diag(np.sqrt(np.abs(w[-rK:])))).reshape(pdim, pdim, rK),
        )
        # Trace normalization of Choi Matrix
        K[i] = K[i] / np.sqrt(np.einsum("ijk,ijk", K[i], K[i].conj())) * np.sqrt(pdim)
    return np.array(K)


def sampled_measurements(y: np.ndarray, n_samples: int) -> np.ndarray:
    """Compute finite sample estimates of input probabilities.

    Args:
        y: 2D array of measurement outcomes for different sequences, Each column contains the outcome probabilities for
            a fixed sequence
        n_samples: Number of samples for each experiment

    Returns:
        y_sampled: 2D array of sampled measurement outcomes for different sequences;
            Each column contains the sampled outcome probabilities for a fixed sequence

    Notes:
        The entries of each column of y form a probability distribution.
        From this distribution n random samples are drawn which give estimates for
        the initial probabilities. This simulates finite sample size data.

    """
    n_povm, m = y.shape
    cap_tolerance = 1e-6
    if any(y.reshape(-1) > 1) or any(y.reshape(-1) < 0):
        y_new = np.maximum(np.minimum(y, 1), 0)
        if np.sum(np.abs(y_new - y)) > cap_tolerance:
            warnings.warn(
                f"Warning: Probabilities capped to interval [0,1], l1-difference to input: {np.sum(np.abs(y_new - y))}",
                stacklevel=2,
            )
        y = y_new
    rng = np.random.default_rng()
    y_sampled = np.zeros(y.shape)
    for i in range(m):
        y_sampled[:, i] = rng.multinomial(n_samples, [y[o, i] for o in range(n_povm)]) / n_samples
    return y_sampled


def random_len_seq(d: int, min_l: int, max_l: int, N: int) -> np.ndarray:
    """Generate random gate sequence instructions.

    These contain sequences of different lengths,
    the lengths are drawn uniformly at random from (min_l, ..., max_l).

    Args:
        d: Number of gates
        min_l: Minimum sequence length
        max_l: Maximum sequence length
        N: Number of random sequences

    Returns:
        J: 2D array where each row contains the gate indices of a gate sequence

    """
    seq_lengths = np.random.randint(min_l, max_l + 1, N)
    J = []
    for length in seq_lengths:
        j_curr = np.random.randint(0, d, length)
        J.append(list(np.pad(j_curr, (0, max_l - length), "constant", constant_values=-1)))
    return np.array(J)


def generate_fids(d: int, length: int, m_f: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Generate random fiducial sequencecs.

    Args:
        d: Number of gates
        length: Total sequence length
        m_f: Number of random fiducial sequences

    Returns:
        J_fid: Sequence list of only the fiducial sequences
        J_fid2: Sequence list for all combinations of two concatenated fiducial sequences
        J_meas: Sequence list for all combinations of fiducials seuqneces with a gate in between:
            fiducial1 -- gate -- fiducial2

    """
    fid_length = (length - 1) // 2
    fid = random.sample(range(d**fid_length), m_f).copy()
    J_fid = [list(local_basis(ind, d, fid_length)) for ind in fid]
    J_fid2 = np.array([seqL + seqR for seqL in J_fid for seqR in J_fid])
    J_meas = np.zeros((d, m_f**2, length), dtype="int")
    for k in range(d):
        J_meas[k] = np.array([seqL + [k] + seqR for seqL in J_fid for seqR in J_fid])
    return np.array(J_fid), J_fid2, J_meas


def tvd(X: np.ndarray, E: np.ndarray, rho: np.ndarray, J: np.ndarray, y_data: np.ndarray) -> float:
    """Return the total variation distance between model probabilities and y_data.

    Args:
        X: Gate set
        E: POVM
        rho: Initial state
        J: 2D array where each row contains the gate indices of a gate sequence
        y_data: 2D array of measurement outcomes for sequences in J;
            Each column contains the outcome probabilities for a fixed sequence

    Returns:
        dist: The total variation distance.

    """
    n_povm = y_data.shape[0]
    y_model = np.real(np.array([[E[i].conj() @ contract(X, j) @ rho for j in J] for i in range(n_povm)]))
    dist = float(la.norm(y_model - y_data, ord=1)) / 2
    return dist


def random_seq_design(d: int, l_min: int, l_cut: int, l_max: int, N_short: int, N_long: int) -> np.ndarray:
    """Generate a set of random sequences ith given lengths.

    Disclaimer: Draws without replacement but inefficiently (not working for sequence length > 24)
    This sequence lengths (circuit depths) are chosen for a mix of very short sequences (better convergence) and some
    slightly longer sequences to reduce the generalization error.
    This sequence design is heuristic and intended for coarse and fast estimates. For very accurate estimates at the
    cost of higher measurement effort it is recommended to use pyGSTi with long sequence GST.

    Args:
        d: The number of gates in the gate set
        l_min: Minimum sequence lenght
        l_cut: Cutoff sequence lenght: N_short sequences are equally distributed among lengths l_min < l < l_cut
        l_max: N_long sequences are equally distributed among lengths l_cut + 1 < l < l_max. Currently l_max < 24 only.
        N_short: Number of short sequences
        N_long: Number of long sequences

    Returns:
        J: 2D array where each row contains the gate indices of a gate sequence

    """
    # Open problems:
    # - Change randomness to work with longer sequences;
    # - Handle case where Number of sequences is smaller than the available lengths

    max_supported = 23
    if l_max > max_supported:
        raise ValueError(f"Currently only sequences lenghts <= {max_supported} are supported.")

    J = np.array([-np.ones(l_max)])
    # Short sequences:
    seq_counts = []
    N_remaining = N_short
    for length in range(l_min, l_cut + 1):
        seq_counts.append(int(np.min([np.floor(N_remaining / (l_cut + 1 - length)), d**length])))
        if seq_counts[-1] > 0:
            ind_curr = np.array(random.sample(range(d**length), seq_counts[-1]))
            J_curr = np.array(
                [
                    np.pad(
                        local_basis(ind, d, length),
                        (0, l_max - length),
                        "constant",
                        constant_values=-1,
                    )
                    for ind in ind_curr
                ]
            )
            J = np.concatenate((J, J_curr), axis=0)
            N_remaining = N_short - 1 - np.sum(seq_counts)
    if N_remaining > 0:
        print(
            "Number of possible sequences without replacement for the given sequence\
        length range is lower than the desired total number of sequences"
        )

    # Long sequences:
    seq_counts = []
    N_remaining = N_long
    for length in range(l_cut + 1, l_max + 1):
        seq_counts.append(int(np.min([np.floor(N_remaining / (l_max + 1 - length)), d**length])))
        if seq_counts[-1] > 0:
            ind_curr = np.array(random.sample(range(d**length), seq_counts[-1]))
            J_curr = np.array(
                [
                    np.pad(
                        local_basis(ind, d, length),
                        (0, l_max - length),
                        "constant",
                        constant_values=-1,
                    )
                    for ind in ind_curr
                ]
            )
            J = np.concatenate((J, J_curr), axis=0)
            N_remaining = N_long - np.sum(seq_counts)
    return J.astype(int)


def n_params(pdim: int, d: int, rK: int, n_povm: int) -> int:
    """Returns the number of free parameters in a gate set.

    Unitary gauge parameters and unitary freedom in representing the Kraus operators and the POVM elements are
    subtracted. For details, see
    https://docserv.uni-duesseldorf.de/servlets/DerivateServlet/Derivate-71735/dissertation_brieger-2%281%29.pdf

    Args:
        pdim: The physical dimension, i.e. 2^(#qubits)
        d: The number of gates in the gate set
        rK: The Kraus rank of the gates
        n_povm: The number of POVM elements; For computational basis measurements n_povm = pdim

    Returns:
        n_params: The number of free parameters

    """
    # Order: gates + stat + povm - povm_freedom - gauge freedom - Kraus freedom
    return d * (pdim**2 * (2 * rK - 1) - rK**2) + pdim * (n_povm * pdim - pdim)


def multikron(matrix_array: np.ndarray) -> np.ndarray:
    """Computes the Kronecker product of all matrices in an array.

    The order is matrix_array[0] otimes matrix_array[1] otimes ...

    Args:
        matrix_array: An array containing matrices

    Returns:
        res: The resulting tensor product (potentially a very large array)

    """
    res = matrix_array[0]
    for i in range(1, matrix_array.shape[0]):
        res = np.kron(res, matrix_array[i])
    return res


def local_dephasing_pp(prob_vec: list[float]) -> np.ndarray:
    """Returns the tensor product of single qubit dephasing channels in Pauli basis.

    Args:
        prob_vec: A list of dephasing probabilities

    Returns:
        D_final: Process matrix of the tensor product of local dephasing channels

    """
    D_loc_array = np.array([[[1, 0, 0, 0], [0, 1 - 2 * p, 0, 0], [0, 0, 1 - 2 * p, 0], [0, 0, 0, 1]] for p in prob_vec])
    D_final = multikron(D_loc_array)
    return D_final


def outcome_probs_from_files(
    folder_name: str, basis_dict: dict[str, int], n_povm: int, N: int
) -> tuple[np.ndarray, int]:
    """Searches a specified folder for .txt files containing circuit outcomes and combines the results.

    Each text file needs to have line in the following format:
    1: 1,0,1,0,1,1,0,0,0
    2: 1,1,1,1,0,1,1,1,1
    Here the first number specifies the circuit and basis outcome of each shot is written to the right.

    Args:
        folder_name: The relative or absolute name of the data folder
        basis_dict: Translation between label for each shot in the .txt files and the numbering of the POVM element.
            Example (two qubits): {"00": 0, "01": 1, "10": 2, "11": 3,}
        n_povm: The number of POVM elements in the data set
        N: The number of circuits in the data set

    Returns:
        y: 2D array of measurement outcomes for all measured sequences; Each column contains the outcome probabilities
            for a fixed sequence
        avg_counts: The average number of shots per circuit after combining data sets

    """
    filenames = os.listdir(path=folder_name)
    datafile_names = [s for s in filenames if ".txt" in s]

    # array of outcome probabilities:
    y = np.zeros((n_povm, N))
    sample_counts = np.zeros((len(datafile_names), N))
    k = 0
    for filename in datafile_names:
        with open(folder_name + "/" + filename, encoding="ASCII") as file:
            i = 0
            for line in file:
                # removing row index at beginning and \n mark at end
                line_entries = line.rstrip().split(": ")
                # splitting result string at commas
                result_list = line_entries[1:][0].split(",")
                # translate each measurement result onto basis index
                for entry in result_list:
                    j = basis_dict[entry]
                    y[j, i] += 1
                sample_counts[k, i] = len(result_list)
                i += 1
        k += 1
    total_counts = np.sum(sample_counts, axis=0)
    avg_counts = int(np.average(total_counts))
    return y / total_counts, avg_counts
