# -*- coding: utf-8 -*-

# Copyright 2019, IBM.
#
# This source code is licensed under the Apache License, Version 2.0 found in
# the LICENSE.txt file in the root directory of this source tree.

# pylint: disable=unused-argument,too-many-return-statements,len-as-condition
"""
Transformations between QuantumChannel represenations.
"""

import numpy as np
import scipy.linalg as la

from qiskit.qiskiterror import QiskitError
from qiskit.quantum_info.operators.predicates import is_hermitian_matrix
from qiskit.quantum_info.operators.predicates import ATOL_DEFAULT


def _to_choi(rep, data, input_dim, output_dim):
    """Transform a QuantumChannel to the Choi representation."""
    if rep == 'Choi':
        return data
    if rep == 'Operator':
        return _from_operator('Choi', data, input_dim, output_dim)
    if rep == 'SuperOp':
        return _superop_to_choi(data, input_dim, output_dim)
    if rep == 'Kraus':
        return _kraus_to_choi(data, input_dim, output_dim)
    if rep == 'Chi':
        return _chi_to_choi(data, input_dim, output_dim)
    if rep == 'PTM':
        data = _ptm_to_superop(data, input_dim, output_dim)
        return _superop_to_choi(data, input_dim, output_dim)
    if rep == 'Stinespring':
        return _stinespring_to_choi(data, input_dim, output_dim)
    raise QiskitError('Invalid QuantumChannel {}'.format(rep))


def _to_superop(rep, data, input_dim, output_dim):
    """Transform a QuantumChannel to the SuperOp representation."""
    if rep == 'SuperOp':
        return data
    if rep == 'Operator':
        return _from_operator('SuperOp', data, input_dim, output_dim)
    if rep == 'Choi':
        return _choi_to_superop(data, input_dim, output_dim)
    if rep == 'Kraus':
        return _kraus_to_superop(data, input_dim, output_dim)
    if rep == 'Chi':
        data = _chi_to_choi(data, input_dim, output_dim)
        return _choi_to_superop(data, input_dim, output_dim)
    if rep == 'PTM':
        return _ptm_to_superop(data, input_dim, output_dim)
    if rep == 'Stinespring':
        return _stinespring_to_superop(data, input_dim, output_dim)
    raise QiskitError('Invalid QuantumChannel {}'.format(rep))


def _to_kraus(rep, data, input_dim, output_dim):
    """Transform a QuantumChannel to the Kraus representation."""
    if rep == 'Kraus':
        return data
    if rep == 'Stinespring':
        return _stinespring_to_kraus(data, input_dim, output_dim)
    if rep == 'Operator':
        return _from_operator('Kraus', data, input_dim, output_dim)
    # Convert via Choi and Kraus
    if rep != 'Choi':
        data = _to_choi(rep, data, input_dim, output_dim)
    return _choi_to_kraus(data, input_dim, output_dim)


def _to_chi(rep, data, input_dim, output_dim):
    """Transform a QuantumChannel to the Chi representation."""
    if rep == 'Chi':
        return data
    # Check valid n-qubit input
    _check_nqubit_dim(input_dim, output_dim)
    if rep == 'Operator':
        return _from_operator('Chi', data, input_dim, output_dim)
    # Convert via Choi representation
    if rep != 'Choi':
        data = _to_choi(rep, data, input_dim, output_dim)
    return _choi_to_chi(data, input_dim, output_dim)


def _to_ptm(rep, data, input_dim, output_dim):
    """Transform a QuantumChannel to the PTM representation."""
    if rep == 'PTM':
        return data
    # Check valid n-qubit input
    _check_nqubit_dim(input_dim, output_dim)
    if rep == 'Operator':
        return _from_operator('PTM', data, input_dim, output_dim)
    # Convert via Superoperator representation
    if rep != 'SuperOp':
        data = _to_superop(rep, data, input_dim, output_dim)
    return _superop_to_ptm(data, input_dim, output_dim)


def _to_stinespring(rep, data, input_dim, output_dim):
    """Transform a QuantumChannel to the Stinespring representation."""
    if rep == 'Stinespring':
        return data
    if rep == 'Operator':
        return _from_operator('Stinespring', data, input_dim, output_dim)
    # Convert via Superoperator representation
    if rep != 'Kraus':
        data = _to_kraus(rep, data, input_dim, output_dim)
    return _kraus_to_stinespring(data, input_dim, output_dim)


def _to_operator(rep, data, input_dim, output_dim):
    """Transform a QuantumChannel to the Operator representation."""
    if rep == 'Operator':
        return data
    if rep == 'Stinespring':
        return _stinespring_to_operator(data, input_dim, output_dim)
    # Convert via Kraus representation
    if rep != 'Kraus':
        data = _to_kraus(rep, data, input_dim, output_dim)
    return _kraus_to_operator(data, input_dim, output_dim)


def _from_operator(rep, data, input_dim, output_dim):
    """Transform Operator representation to other representation."""
    if rep == 'Operator':
        return data
    if rep == 'SuperOp':
        return np.kron(np.conj(data), data)
    if rep == 'Choi':
        vec = np.ravel(data, order='F')
        return np.outer(vec, np.conj(vec))
    if rep == 'Kraus':
        return ([data], None)
    if rep == 'Stinespring':
        return (data, None)
    if rep == 'Chi':
        _check_nqubit_dim(input_dim, output_dim)
        data = _from_operator('Choi', data, input_dim, output_dim)
        return _choi_to_chi(data, input_dim, output_dim)
    if rep == 'PTM':
        _check_nqubit_dim(input_dim, output_dim)
        data = _from_operator('SuperOp', data, input_dim, output_dim)
        return _superop_to_ptm(data, input_dim, output_dim)
    raise QiskitError('Invalid QuantumChannel {}'.format(rep))


def _kraus_to_operator(data, input_dim, output_dim):
    """Transform Kraus representation to Operator representation."""
    if data[1] is not None or len(data[0]) > 1:
        raise QiskitError(
            'Channel cannot be converted to Operator representation')
    return data[0][0]


def _stinespring_to_operator(data, input_dim, output_dim):
    """Transform Stinespring representation to Operator representation."""
    trace_dim = data[0].shape[0] // output_dim
    if data[1] is not None or trace_dim != 1:
        raise QiskitError(
            'Channel cannot be converted to Operator representation')
    return data[0]


def _superop_to_choi(data, input_dim, output_dim):
    """Transform SuperOp representation to Choi representation."""
    shape = (output_dim, output_dim, input_dim, input_dim)
    return _reshuffle(data, shape)


def _choi_to_superop(data, input_dim, output_dim):
    """Transform Choi to SuperOp representation."""
    shape = (input_dim, output_dim, input_dim, output_dim)
    return _reshuffle(data, shape)


def _kraus_to_choi(data, input_dim, output_dim):
    """Transform Kraus representation to Choi representation."""
    choi = 0
    kraus_l, kraus_r = data
    if kraus_r is None:
        for i in kraus_l:
            vec = i.ravel(order='F')
            choi += np.outer(vec, vec.conj())
    else:
        for i, j in zip(kraus_l, kraus_r):
            choi += np.outer(i.ravel(order='F'), j.ravel(order='F').conj())
    return choi


def _choi_to_kraus(data, input_dim, output_dim, atol=ATOL_DEFAULT):
    """Transform Choi representation to Kraus representation."""
    # Check if hermitian matrix
    if is_hermitian_matrix(data, atol=atol):
        # Get eigen-decomposition of Choi-matrix
        w, v = la.eigh(data)
        # Check eigenvaleus are non-negative
        if len(w[w < -atol]) == 0:
            # CP-map Kraus representation
            kraus = []
            for val, vec in zip(w, v.T):
                if abs(val) > atol:
                    k = np.sqrt(val) * vec.reshape(
                        (output_dim, input_dim), order='F')
                    kraus.append(k)
            # If we are converting a zero matrix, we need to return a Kraus set
            # with a single zero-element Kraus matrix
            if not kraus:
                kraus.append(np.zeros((output_dim, input_dim), dtype=complex))
            return (kraus, None)
    # Non-CP-map generalized Kraus representation
    mat_u, svals, mat_vh = la.svd(data)
    kraus_l = []
    kraus_r = []
    for val, vec_l, vec_r in zip(svals, mat_u.T, mat_vh.conj()):
        kraus_l.append(
            np.sqrt(val) * vec_l.reshape((output_dim, input_dim), order='F'))
        kraus_r.append(
            np.sqrt(val) * vec_r.reshape((output_dim, input_dim), order='F'))
    return (kraus_l, kraus_r)


def _stinespring_to_kraus(data, input_dim, output_dim):
    """Transform Stinespring representation to Kraus representation."""
    kraus_pair = []
    for stine in data:
        if stine is None:
            kraus_pair.append(None)
        else:
            trace_dim = stine.shape[0] // output_dim
            iden = np.eye(output_dim)
            kraus = []
            for j in range(trace_dim):
                vec = np.zeros(trace_dim)
                vec[j] = 1
                kraus.append(np.kron(iden, vec[None, :]).dot(stine))
            kraus_pair.append(kraus)
    return tuple(kraus_pair)


def _stinespring_to_choi(data, input_dim, output_dim):
    """Transform Stinespring representation to Choi representation."""
    trace_dim = data[0].shape[0] // output_dim
    stine_l = np.reshape(data[0], (output_dim, trace_dim, input_dim))
    if data[1] is None:
        stine_r = stine_l
    else:
        stine_r = np.reshape(data[1], (output_dim, trace_dim, input_dim))
    return np.reshape(
        np.einsum('iAj,kAl->jilk', stine_l, stine_r.conj()),
        2 * [input_dim * output_dim])


def _stinespring_to_superop(data, input_dim, output_dim):
    """Transform Stinespring representation to SuperOp representation."""
    trace_dim = data[0].shape[0] // output_dim
    stine_l = np.reshape(data[0], (output_dim, trace_dim, input_dim))
    if data[1] is None:
        stine_r = stine_l
    else:
        stine_r = np.reshape(data[1], (output_dim, trace_dim, input_dim))
    return np.reshape(
        np.einsum('iAj,kAl->ikjl', stine_r.conj(), stine_l),
        (output_dim * output_dim, input_dim * input_dim))


def _kraus_to_stinespring(data, input_dim, output_dim):
    """Transform Kraus representation to Stinespring representation."""
    stine_pair = [None, None]
    for i, kraus in enumerate(data):
        if kraus is not None:
            num_kraus = len(kraus)
            stine = np.zeros((output_dim * num_kraus, input_dim),
                             dtype=complex)
            for j, mat in enumerate(kraus):
                vec = np.zeros(num_kraus)
                vec[j] = 1
                stine += np.kron(mat, vec[:, None])
            stine_pair[i] = stine
    return tuple(stine_pair)


def _kraus_to_superop(data, input_dim, output_dim):
    """Transform Kraus representation to SuperOp representation."""
    kraus_l, kraus_r = data
    superop = 0
    if kraus_r is None:
        for i in kraus_l:
            superop += np.kron(np.conj(i), i)
    else:
        for i, j in zip(kraus_l, kraus_r):
            superop += np.kron(np.conj(j), i)
    return superop


def _chi_to_choi(data, input_dim, output_dim):
    """Transform Chi representation to a Choi representation."""
    num_qubits = int(np.log2(input_dim))
    return _transform_from_pauli(data, num_qubits)


def _choi_to_chi(data, input_dim, output_dim):
    """Transform Choi representation to the Chi representation."""
    num_qubits = int(np.log2(input_dim))
    return _transform_to_pauli(data, num_qubits)


def _ptm_to_superop(data, input_dim, output_dim):
    """Transform PTM representation to SuperOp representation."""
    num_qubits = int(np.log2(input_dim))
    return _transform_from_pauli(data, num_qubits)


def _superop_to_ptm(data, input_dim, output_dim):
    """Transform SuperOp representation to PTM representation."""
    num_qubits = int(np.log2(input_dim))
    return _transform_to_pauli(data, num_qubits)


def _bipartite_tensor(mat1, mat2, shape1=None, shape2=None):
    """Tensor product (A ⊗ B) to bipartite matrices and reravel indicies.

    This is used for tensor product of superoperators and Choi matrices.

    Args:
        mat1 (matrix_like): a bipartite matrix A
        mat2 (matrix_like): a bipartite matrix B
        shape1 (tuple): bipartite-shape for matrix A (a0, a1, a2, a3)
        shape2 (tuple): bipartite-shape for matrix B (b0, b1, b2, b3)

    Returns:
        np.array: a bipartite matrix for reravel(A ⊗ B).

    Raises:
        QiskitError: if input matrices are wrong shape.
    """
    # Convert inputs to numpy arrays
    mat1 = np.array(mat1)
    mat2 = np.array(mat2)

    # Determine bipartite dimensions if not provided
    dim_a0, dim_a1 = mat1.shape
    dim_b0, dim_b1 = mat2.shape
    if shape1 is None:
        sdim_a0 = int(np.sqrt(dim_a0))
        sdim_a1 = int(np.sqrt(dim_a1))
        shape1 = (sdim_a0, sdim_a0, sdim_a1, sdim_a1)
    if shape2 is None:
        sdim_b0 = int(np.sqrt(dim_b0))
        sdim_b1 = int(np.sqrt(dim_b1))
        shape2 = (sdim_b0, sdim_b0, sdim_b1, sdim_b1)
    # Check dimensions
    if len(shape1) != 4 or shape1[0] * shape1[1] != dim_a0 or \
            shape1[2] * shape1[3] != dim_a1:
        raise QiskitError("Invalid shape_a")
    if len(shape2) != 4 or shape2[0] * shape2[1] != dim_b0 or \
            shape2[2] * shape2[3] != dim_b1:
        raise QiskitError("Invalid shape_b")

    return _reravel(mat1, mat2, shape1, shape2)


def _reravel(mat1, mat2, shape1, shape2):
    """Reravel two bipartite matrices."""
    # Reshuffle indicies
    left_dims = shape1[:2] + shape2[:2]
    right_dims = shape1[2:] + shape2[2:]
    tensor_shape = left_dims + right_dims
    final_shape = (np.product(left_dims), np.product(right_dims))
    # Tensor product matrices
    data = np.kron(mat1, mat2)
    data = np.reshape(
        np.transpose(np.reshape(data, tensor_shape), (0, 2, 1, 3, 4, 6, 5, 7)),
        final_shape)
    return data


def _transform_to_pauli(data, num_qubits):
    """Change of basis of bipartite matrix represenation."""
    # Change basis: um_{i=0}^3 |i>><\sigma_i|
    basis_mat = np.array(
        [[1, 0, 0, 1], [0, 1, 1, 0], [0, -1j, 1j, 0], [1, 0j, 0, -1]],
        dtype=complex)
    # Note that we manually renormalized after change of basis
    # to avoid rounding errors from square-roots of 2.
    cob = basis_mat
    for _ in range(num_qubits - 1):
        dim = int(np.sqrt(len(cob)))
        cob = np.reshape(
            np.transpose(
                np.reshape(
                    np.kron(basis_mat, cob), (4, dim * dim, 2, 2, dim, dim)),
                (0, 1, 2, 4, 3, 5)), (4 * dim * dim, 4 * dim * dim))
    return np.dot(np.dot(cob, data), cob.conj().T) / 2**num_qubits


def _transform_from_pauli(data, num_qubits):
    """Change of basis of bipartite matrix represenation."""
    # Change basis: sum_{i=0}^3 =|\sigma_i>><i|
    basis_mat = np.array(
        [[1, 0, 0, 1], [0, 1, 1j, 0], [0, 1, -1j, 0], [1, 0j, 0, -1]],
        dtype=complex)
    # Note that we manually renormalized after change of basis
    # to avoid rounding errors from square-roots of 2.
    cob = basis_mat
    for _ in range(num_qubits - 1):
        dim = int(np.sqrt(len(cob)))
        cob = np.reshape(
            np.transpose(
                np.reshape(
                    np.kron(basis_mat, cob), (2, 2, dim, dim, 4, dim * dim)),
                (0, 2, 1, 3, 4, 5)), (4 * dim * dim, 4 * dim * dim))
    return np.dot(np.dot(cob, data), cob.conj().T) / 2**num_qubits


def _reshuffle(mat, shape):
    """Reshuffle the indicies of a bipartite matrix A[ij,kl] -> A[lj,ki]."""
    return np.reshape(
        np.transpose(np.reshape(mat, shape), (3, 1, 2, 0)),
        (shape[3] * shape[1], shape[0] * shape[2]))


def _check_nqubit_dim(input_dim, output_dim):
    """Return true if dims correspond to an n-qubit channel."""
    if input_dim != output_dim:
        raise QiskitError(
            'Not an n-qubit channel: input_dim' +
            ' ({}) != output_dim ({})'.format(input_dim, output_dim))
    num_qubits = int(np.log2(input_dim))
    if 2**num_qubits != input_dim:
        raise QiskitError('Not an n-qubit channel: input_dim != 2 ** n')
