#               This file is part of the QuTIpy package.
#                https://github.com/sumeetkhatri/QuTIpy
#
#                   Copyright (c) 2022 Sumeet Khatri.
#                       --.- ..- - .. .--. -.--
#
#
# SPDX-License-Identifier: AGPL-3.0
#
#  This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, version 3.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#

import itertools

import cvxpy as cvx
import numpy as np
from numpy.linalg import matrix_power
from scipy.linalg import eig

from qutipy.general_functions import (
    Tr,
    dag,
    eye,
    ket,
    partial_trace,
    syspermute,
    tensor,
)
from qutipy.linalg import gram_schmidt
from qutipy.misc import cvxpy_to_numpy, numpy_to_cvxpy
from qutipy.pauli import (
    generate_nQubit_Pauli,
    generate_nQubit_Pauli_X,
    generate_nQubit_Pauli_Z,
)
from qutipy.states import MaxEnt_state, RandomStateVector
from qutipy.weyl import discrete_Weyl_Z


def Choi_to_Natural(C_AB, dimA, dimB):
    """
    Takes the Choi representation of a map and outputs its natural representation.

    The Choi represenatation Q of the channel acts as:

        vec(N(rho))=Q*vec(rho),

    where N is the channel in question. It can be obtained from the Choi representation
    with a simple reshuffling of indices.
    """

    C_AB = np.array(C_AB)

    return np.array(
        np.reshape(C_AB, [dimA, dimB, dimA, dimB])
        .transpose((0, 2, 1, 3))
        .reshape([dimA * dimA, dimB * dimB])
    ).T


def bit_flip_channel(p):
    """
    Generates the channel rho -> (1-p)*rho+p*X*rho*X.
    """

    return Pauli_channel(p, 0, 0)


def completely_dephasing_channel(d):
    """
    Generates the completely dephasing channel in d dimensions. This channel
    eliminates the off-diagonal elements (in the standard basis) of the input operator.
    """

    if d == 2:
        p = 1 / 2
        return dephasing_channel(p, d=d)[0]
    else:
        p = (1 / d) * np.ones(d)
        return dephasing_channel(p, d=d)


def Kraus_representation(P, dimA, dimB):
    """
    Takes a Choi representation P of a channel and returns its Kraus representation.

    The Choi representation is defined with the channel acting on the second half of
    the maximally entangled vector.
    """

    D, U = eig(P)

    U_cols = U.shape[1]

    # Need to check if the matrix U generated by eig is unitary (up to
    # numerical precision)
    check1 = np.allclose(eye(dimA * dimB), U @ dag(U))
    check2 = np.allclose(eye(dimA * dimB), dag(U) @ U)

    if check1 and check2:
        U = np.array(U)

    # If U is not unitary, use Gram-Schmidt to make it unitary (i.e., make the
    # columns of U orthonormal)
    else:
        C = gram_schmidt([U[:, i] for i in range(U_cols)], dimA * dimB)
        U = np.sum([tensor(dag(ket(U_cols, i)), C[i]) for i in range(U_cols)], 0)
        # print(U)
    K = []

    for i in range(U_cols):
        Col = U[:, i]
        K_tmp = np.array(np.sqrt(D[i]) * Col.reshape([dimA, dimB]))
        K.append(K_tmp.transpose())

    return K


def phase_damping_channel(p):
    """
    Generates the phase damping channel.
    """

    K1 = np.array([[1, 0], [0, np.sqrt(p)]])
    K2 = np.array([[0, 0], [0, np.sqrt(1 - p)]])

    return [K1, K2]


def generate_channel_isometry(K, dimA, dimB):
    """
    Generates an isometric extension of the
    channel specified by the Kraus operators K. dimA is the dimension of the
    input space of the channel, and dimB is the dimension of the output space
    of the channel. If dimA=dimB, then the function also outputs a unitary
    extension of the channel given by a particular construction.
    """

    dimE = len(K)

    V = np.sum([tensor(K[i], ket(dimE, i)) for i in range(dimE)], 0)

    if dimA == dimB:
        # In this case, the unitary we generate has dimensions dimA*dimE x
        # dimA*dimE
        U = tensor(V, dag(ket(dimE, 0)))
        states = [V @ ket(dimA, i) for i in range(dimA)]
        for i in range(dimA * dimE - dimA):
            states.append(RandomStateVector(dimA * dimE))

        states_new = gram_schmidt(states, dimA * dimE)

        count = dimA
        for i in range(dimA):
            for j in range(1, dimE):
                U = U + tensor(states_new[count], dag(ket(dimA, i)), dag(ket(dimE, j)))
                count += 1

        return V, np.array(U)
    else:
        return V


def Pauli_channel_nQubit(n, p, alt_repr=False):
    """
    Generates the Kraus operators, an isometric extension, and a unitary
    extension of the n-qubit Pauli channel specified by the 2^(2*n) parameters in
    p, which must be probabilities in order for the map to be a channel. (i.e.,
    they must be non-negative and sum to one.)

    If alt_repr=True, then the channel is of the form

    P(rho)=\\sum_{a,b} p_{a,b} X^aZ^b(rho)Z^bX^a

    where a and b are n-bit strings
    (using the n-qubit X and Z operators as generated by the functions
    generate_nQubit_Pauli_X and generate_nQubit_Pauli_Z).
    """

    K = []

    if not alt_repr:
        S = list(itertools.product(*[range(0, 4)] * n))
        for i in range(2 ** (2 * n)):
            K.append(np.sqrt(p[i]) * generate_nQubit_Pauli(list(S[i])))

        V, U = generate_channel_isometry(K, 2**n, 2**n)

        return K, V, U

    else:  # alt_repr==True
        S = list(itertools.product(*[range(0, 2)] * n))
        count = 0
        for a in S:
            a = list(a)
            for b in S:
                b = list(b)
                K.append(
                    np.sqrt(p[count])
                    * generate_nQubit_Pauli_X(a)
                    @ generate_nQubit_Pauli_Z(b)
                )
                count += 1

        V, U = generate_channel_isometry(K, 2**n, 2**n)

        return K, V, U


def apply_channel(K, rho, sys=None, dim=None, adjoint=False):
    """
    Applies the channel with Kraus operators in K to the state rho on
    systems specified by sys. The dimensions of the subsystems on which rho
    acts are given by dim.

    If adjoint is True, then this function applies the adjoint of the given
    channel.
    """

    if isinstance(rho, cvx.Variable):
        rho = cvxpy_to_numpy(rho)
        rho_out = apply_channel(K, rho, sys, dim, adjoint)
        return numpy_to_cvxpy(rho_out)

    if adjoint:
        K_tmp = K
        K = []
        K = [dag(K_tmp[i]) for i in range(len(K_tmp))]

    if sys is None:
        return np.sum([K[i] @ rho @ dag(K[i]) for i in range(len(K))], 0)
    else:
        A = []
        for i in range(len(K)):
            X = 1
            for j in range(len(dim)):
                if j + 1 == sys:
                    X = tensor(X, K[i])
                else:
                    X = tensor(X, eye(dim[j]))
            A.append(X)

        return np.sum([A[i] @ rho @ dag(A[i]) for i in range(len(A))], 0)


def amplitude_damping_channel(gamma):
    """
    Generates the amplitude damping channel.
    """

    A1 = np.array([[1, 0], [0, np.sqrt(1 - gamma)]])
    A2 = np.array([[0, np.sqrt(gamma)], [0, 0]])

    return [A1, A2]


def Natural_representation(K):
    """
    Calculates the natural representation of the channel (in the standard basis)
    given by the Kraus operators in K. In terms of the Kraus operators, the natural
    representation of the channel in the standard basis is given by

    N=sum_i K_i ⊗ conj(K_i),

    where the sum is over the Kraus operators K_i in K.
    """

    return np.sum([tensor(k, np.conjugate(k)) for k in K], 0)


def BB84_channel(Q):
    """
    Generates the channel corresponding to the BB84 protocol with
    equal X and Z errors, given by the QBER Q. The definition of this
    channel can be found in:

        "Additive extensions of a quantum channel", by
            Graeme Smith and John Smolin. (arXiv:0712.2471)

    """

    return Pauli_channel(Q - Q**2, Q**2, Q - Q**2)


def Choi_representation(K, dimA):
    """
    Calculates the Choi representation of the map with Kraus operators K.
    dimA is the dimension of the input space of the channel.

    The Choi represenatation is defined with the channel acting on the second
    half of the maximally entangled vector.
    """

    Gamma = MaxEnt_state(dimA, normalized=False)

    return np.array(apply_channel(K, Gamma, 2, [dimA, dimA]), dtype=np.complex_)


def compose_channels(C):
    """
    Takes a composition of channels. The variable C should be a list of lists,
    with each list consisting of the Kraus operators of the channels to be composed.

    If C=[K1,K2,...,Kn], then this function returns the composition such that
    the channel corresponding to K1 is applied first, then K2, etc.
    """

    d = C[0][0].shape[0]

    lengths = []
    for c in C:
        lengths.append(len(c))

    combs = list(itertools.product(*[range(length) for length in lengths]))

    K_n = []

    for comb in combs:
        # tmp=1
        tmp = eye(d)
        for i in range(len(comb)):
            tmp = C[i][comb[i]] @ tmp
        K_n.append(tmp)

    return K_n


def tensor_channels(C):
    """
    Takes the tensor product of the channels in C.

    C is a set of sets of Kraus operators.
    """

    lengths = []
    for c in C:
        lengths.append(len(c))

    combs = list(itertools.product(*[range(length) for length in lengths]))

    K_n = []

    for comb in combs:
        tmp = 1
        for i in range(len(comb)):
            tmp = tensor(tmp, C[i][comb[i]])
        K_n.append(tmp)

    return K_n


def depolarizing_channel_n_uses(p, n, rho, m):
    """
    Generates the output state corresponding to the depolarizing channel
    applied to each one of n systems in the joint state rho. p is the
    depolarizing probability as defined in the function "depolarizing_channel"
    above.

    If rho contains m>n systems, then the first m-n systems are left alone.
    """

    dims = 2 * np.ones(m).astype(int)

    rho_out = np.zeros((2**m, 2**m))

    for k in range(n + 1):
        indices = list(itertools.combinations(range(1, n + 1), k))

        # print k,indices

        for index in indices:
            index = list(index)

            index = np.array(index) + (m - n)
            index = list(index.astype(int))

            index_diff = np.setdiff1d(range(1, m + 1), index)

            perm_arrange = np.append(index, index_diff).astype(int)
            perm_rearrange = np.zeros(m)

            for i in range(m):
                perm_rearrange[i] = np.argwhere(perm_arrange == i + 1)[0][0] + 1

            perm_rearrange = perm_rearrange.astype(int)

            mix = matrix_power(eye(2**k) / 2, k)

            rho_part = partial_trace(rho, index, dims)

            rho_out = rho_out + (4 * p / 3.0) ** k * (1 - (4 * p / 3.0)) ** (
                n - k
            ) * syspermute(tensor(mix, rho_part), perm_rearrange, dims)

    return rho_out


def diamond_norm(J, dimA, dimB, display=False):
    """
    Computes the diamond norm of a superoperator with Choi representation J.
    dimA is the dimension of the input space of the channel, and dimB is the
    dimension of the output space.

    The form of the SDP used comes from Theorem 3.1 of:

        'Simpler semidefinite programs for completely bounded norms',
            Chicago Journal of Theoretical Computer Science 2013,
            by John Watrous
    """

    """
    The Choi representation J in the above paper is defined using a different
    convention:
        J=(N⊗ I)(|Phi^+><Phi^+|).
    In other words, the channel N acts on the first half of the maximally-
    entangled state, while the convention used throughout this code stack
    is
        J=(I⊗ N)(|Phi^+><Phi^+|).
    We thus use syspermute to convert to the form used in the aforementioned
    paper.
    """

    J = syspermute(J, [2, 1], [dimA, dimB])

    X = cvx.Variable((dimA * dimB, dimA * dimB), hermitian=False)
    rho0 = cvx.Variable((dimA, dimA), PSD=True)
    rho1 = cvx.Variable((dimA, dimA), PSD=True)

    M = cvx.bmat([[cvx.kron(eye(dimB), rho0), X], [X.H, cvx.kron(eye(dimB), rho1)]])

    c = []
    c += [M >> 0, cvx.trace(rho0) == 1, cvx.trace(rho1) == 1]

    obj = cvx.Maximize(
        (1 / 2) * cvx.real(cvx.trace(dag(J) @ X))
        + (1 / 2) * cvx.real(cvx.trace(J @ X.H))
    )

    prob = cvx.Problem(obj, constraints=c)

    prob.solve(verbose=display, eps=1e-7)

    return prob.value


def depolarizing_channel_nQubits(n, p):
    """
    For 0<=p<=1, this returns the n-qubit Pauli channel given by
    p[0]=1-p, p[i]=p/(2^(2*n)-1) for all i>=1.
    """

    p = [1 - p] + [p / (2 ** (2 * n) - 1) for i in range(2 ** (2 * n) - 1)]

    return Pauli_channel_nQubit(n, p, alt_repr=True)


def dephasing_channel(p, d=2):
    """
    Generates the channel rho -> (1-p)*rho+p*Z*rho*Z. (In the case d=2.)

    For d>=2, we let p be a list of d probabilities, and we use the discrete Weyl-Z
    operators to define the channel.

    For p=1/d, we get the completely dephasing channel.
    """

    if d == 2:
        return Pauli_channel(0, 0, p)
    else:
        K = [np.sqrt(p[k]) * matrix_power(discrete_Weyl_Z(d), k) for k in range(d)]
        return K


def generalized_amplitude_damping_channel(gamma, N):
    """
    Generates the generalized amplitude damping channel.
    """

    if N == 0:
        return amplitude_damping_channel(gamma)
    elif N == 1:
        A1 = np.array([[np.sqrt(1 - gamma), 0], [0, 1]])
        A2 = np.array([[0, 0], [np.sqrt(gamma), 0]])
        return [A1, A2]
    else:
        A1 = np.sqrt(1 - N) * np.array([[1, 0], [0, np.sqrt(1 - gamma)]])
        A2 = np.sqrt(1 - N) * np.array([[0, np.sqrt(gamma)], [0, 0]])
        A3 = np.sqrt(N) * np.array([[np.sqrt(1 - gamma), 0], [0, 1]])
        A4 = np.sqrt(N) * np.array([[0, 0], [np.sqrt(gamma), 0]])

        return [A1, A2, A3, A4]


def n_channel_uses(K, n):
    """
    Given the Kraus operators K of a channel, this function generates the
    Kraus operators corresponding to the n-fold tensor power of the channel.
    dimA is the dimension of the input space, and dimB the dimension of the
    output space.
    """

    r = len(K)  # Number of Kraus operators

    combs = list(itertools.product(*[range(r)] * n))

    K_n = []

    for comb in combs:
        # print comb
        tmp = 1
        for i in range(n):
            tmp = tensor(tmp, K[comb[i]])
        K_n.append(tmp)

    return K_n


def channel_scalar_multiply(K, x):
    """
    Multiplies the channel with Kraus operators in K by the scalar x.
    This means that each Kraus operator is multiplied by sqrt(x)!
    """

    K_new = []

    for i in range(len(K)):
        K_new.append(np.sqrt(x) * K[i])

    return K_new


def Pauli_channel_coeffs(K, n, as_dict=False):
    """
    Generates the coefficients c_{a,b} such that

        P(X^aZ^b)=c_{a,b}X^aZ^b,

    for the channel P with the Kraus operators in K.
    """

    if as_dict:
        c = {}
    else:
        c = []

    S = list(itertools.product(*[range(0, 2)] * n))
    # print(S)

    for a in S:
        for b in S:
            Xa = generate_nQubit_Pauli_X(list(a))
            Zb = generate_nQubit_Pauli_Z(list(b))
            if as_dict:
                c[(a, b)] = (1 / 2**n) * Tr(dag(Xa @ Zb) @ apply_channel(K, Xa @ Zb))
            else:
                c.append((1 / 2**n) * Tr(dag(Xa @ Zb) @ apply_channel(K, Xa @ Zb)))

    return c


def Pauli_channel(px, py, pz):
    """
    Generates the Kraus operators, an isometric extension, and a unitary
    extension of the one-qubit Pauli channel specified by the parameters px, py, pz.
    """

    pI = 1 - px - py - pz

    Sx = np.array([[0, 1], [1, 0]])
    Sy = np.array([[0, -1j], [1j, 0]])
    Sz = np.array([[1, 0], [0, -1]])

    K = [np.sqrt(pI) * eye(2), np.sqrt(px) * Sx, np.sqrt(py) * Sy, np.sqrt(pz) * Sz]

    V, U = generate_channel_isometry(K, 2, 2)

    return K, V, U


def depolarizing_channel(p):
    """
    For 0<=p<=1, this returns the one-qubit Pauli channel given by px=py=pz=p/3.
    """

    return Pauli_channel(p / 3.0, p / 3.0, p / 3.0)
